test(pytorch): move test_import_gpu cases to single_gpu suite

test_import_gpu[transformer_engine] fails on the CPU-only default-runner
that unit-test runs on because TE 2.17 dlopens libcuda.so.1 at import
time; libcuda is only mounted into containers via nvidia-container-cli
when a GPU is present.

The test is correctly gated with @pytest.mark.skipif(not IS_CUDA) — that
gate is about the IMAGE variant (CUDA vs CPU), not about the test
runner's GPU availability. Moving both parametrized cases (flash_attn,
transformer_engine) into single_gpu/ where the workflow runs on
x86-g6xl-runner with --gpus all gives them a real driver mount at
import time.

flash_attn passed on the CPU fleet only because it defers libcuda dlopen
until the first CUDA call — the import itself succeeded. Moving it too
keeps GPU-only import checks in a single suite and avoids relying on
that lazy-loading behavior.

The IS_CUDA skip is preserved as a belt-and-suspenders guard for any
future CPU-variant coverage that reuses the single_gpu suite.

Unit suite retains COMMON_PACKAGES (torch, torchvision, torchaudio,
deepspeed, boto3, requests, yaml, packaging) — none of which need a
GPU driver at import.

Author: Bhanu Teja Goshikonda

test/pytorch/single_gpu/test_imports.py

@@ -0,0 +1,26 @@
+"""Verify GPU-only Python packages import successfully.
+
+TE 2.17 dlopens libcuda.so.1 at import time; libcuda is only mounted into
+containers via nvidia-container-cli when a GPU is present. That means these
+imports must run on the single-GPU fleet (--gpus all), not on the CPU-only
+default-runner used for unit tests. The IS_CUDA skip is kept as a
+belt-and-suspenders guard for CPU-variant images.
+"""
+
+import importlib
+import os
+
+import pytest
+
+IS_CUDA = os.path.isdir("/usr/local/cuda")
+
+GPU_PACKAGES = [
+    "flash_attn",
+    "transformer_engine",
+]
+
+
+@pytest.mark.skipif(not IS_CUDA, reason="CUDA-only package")
+@pytest.mark.parametrize("package", GPU_PACKAGES)
+def test_import_gpu(package):
+    importlib.import_module(package)

test/pytorch/unit/test_imports.py

@@ -1,12 +1,9 @@
 """Verify key Python packages import successfully."""
 
 import importlib
-import os
 
 import pytest
 
-IS_CUDA = os.path.isdir("/usr/local/cuda")
-
 COMMON_PACKAGES = [
     "torch",
     "torchvision",
@@ -18,18 +15,7 @@
     "packaging",
 ]
 
-GPU_PACKAGES = [
-    "flash_attn",
-    "transformer_engine",
-]
-
 
 @pytest.mark.parametrize("package", COMMON_PACKAGES)
 def test_import(package):
     importlib.import_module(package)
-
-
-@pytest.mark.skipif(not IS_CUDA, reason="CUDA-only package")
-@pytest.mark.parametrize("package", GPU_PACKAGES)
-def test_import_gpu(package):
-    importlib.import_module(package)