Revert "[xpu][test] Port 2 test/prototype/test_{parq, quantized_training} UT files to intel XPU" (#3432)

Revert "[xpu][test] Port 2 test/prototype/test_{parq, quantized_training} UT …"

This reverts commit 40c4f44.

Author: metascroy

test/prototype/test_parq.py

@@ -51,11 +51,7 @@
     torch_version_at_least,
 )
 
-_DEVICE = torch.device(
-    torch.accelerator.current_accelerator().type
-    if torch.accelerator.is_available()
-    else "cpu"
-)
+_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
 class M(nn.Module):

test/prototype/test_quantized_training.py

@@ -34,17 +34,11 @@
     quantize_int8_rowwise,
 )
 from torchao.quantization.quant_api import quantize_
-from torchao.utils import get_current_accelerator_device
 
 if common_utils.SEED is None:
     common_utils.SEED = 1234
 
-_DEVICES = (
-    ["cpu"]
-    + (["cuda"] if torch.cuda.is_available() else [])
-    + (["xpu"] if torch.xpu.is_available() else [])
-)
-_DEVICE = get_current_accelerator_device()
+_DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else [])
 
 
 def _reset():
@@ -188,14 +182,12 @@ def test_int8_weight_only_training(self, compile, device):
         ],
     )
     @parametrize("module_swap", [False, True])
-    @pytest.mark.skipif(
-        not torch.accelerator.is_available(), reason="GPU not available"
-    )
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     def test_int8_mixed_precision_training(self, compile, config, module_swap):
         _reset()
         bsize = 64
         embed_dim = 64
-        device = _DEVICE
+        device = "cuda"
 
         linear = nn.Linear(embed_dim, embed_dim, device=device)
         linear_int8mp = copy.deepcopy(linear)
@@ -229,9 +221,7 @@ def snr(ref, actual):
 
     @pytest.mark.skip("Flaky on CI")
     @parametrize("compile", [False, True])
-    @pytest.mark.skipif(
-        not torch.accelerator.is_available(), reason="GPU not available"
-    )
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     def test_bitnet_training(self, compile):
         # reference implementation
         # https://github.com/microsoft/unilm/blob/master/bitnet/The-Era-of-1-bit-LLMs__Training_Tips_Code_FAQ.pdf
@@ -256,7 +246,7 @@ def forward(self, x):
         _reset()
         bsize = 4
         embed_dim = 32
-        device = _DEVICE
+        device = "cuda"
 
         # only use 1 matmul shape to reduce triton autotune time
         model_ref = nn.Sequential(
@@ -352,7 +342,7 @@ def _run_subtest(self, args):
             dropout_p=0,
         )
         torch.manual_seed(42)
-        base_model = Transformer(model_args).to(_DEVICE)
+        base_model = Transformer(model_args).cuda()
         fsdp_model = copy.deepcopy(base_model)
 
         quantize_(base_model.layers, quantize_fn)
@@ -372,7 +362,7 @@ def _run_subtest(self, args):
 
         torch.manual_seed(42 + self.rank + 1)
         for iter_idx in range(5):
-            inp = torch.randint(0, vocab_size, (batch_size, seq_len), device=_DEVICE)
+            inp = torch.randint(0, vocab_size, (batch_size, seq_len), device="cuda")
             fsdp_optim.zero_grad(set_to_none=(iter_idx % 2 == 0))
             fsdp_loss = fsdp_model(inp).sum()
             fsdp_loss.backward()
@@ -397,18 +387,14 @@ def _run_subtest(self, args):
             )
 
     @skip_if_lt_x_gpu(_FSDP_WORLD_SIZE)
-    @pytest.mark.skipif(
-        not torch.accelerator.is_available(), reason="GPU not available"
-    )
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     def test_precompute_bitnet_scale(self):
         from torchao.prototype.quantized_training.bitnet import (
             get_bitnet_scale,
             precompute_bitnet_scale_for_fsdp,
         )
 
-        model = nn.Sequential(nn.Linear(32, 64), nn.GELU(), nn.Linear(64, 32)).to(
-            _DEVICE
-        )
+        model = nn.Sequential(nn.Linear(32, 64), nn.GELU(), nn.Linear(64, 32)).cuda()
         model_fsdp = copy.deepcopy(model)
         quantize_(model_fsdp, bitnet_training())
         fully_shard(model_fsdp)