test(megatron): make VLM distributed tests NPU-compatible

After main split the integration tests into
``tests/test_megatron_engine_vlm_distributed.py`` and renamed the
torchrun script to ``run_megatron_engine_vlm_distributed.py`` (areal-project#1301),
re-apply the same NPU compatibility shims as the previous combined-file
version:

- ``try: import mindspeed.megatron_adaptor`` at the top of both files so
  MindSpeed's adapters land before mbridge transitive imports on NPU.
- Replace ``CUDA_AVAILABLE`` with ``ACCELERATOR_AVAILABLE`` driven by
  ``current_platform.device_type in ("cuda", "npu")`` and update the
  five ``@pytest.mark.skipif`` decorators / two ``device_count()`` checks.
- ``current_platform.synchronize()`` instead of ``torch.cuda.synchronize()``
  in the torchrun cleanup path.

The model path is already centralized in
``areal/utils/testing_utils.py::DENSE_MODEL_PATHS`` (env override via
``VLM_MODEL_PATH``), so no additional path resolution is needed here.

Author: Adiactive

tests/test_megatron_engine_vlm_distributed.py

@@ -17,9 +17,14 @@
 import sys
 
 import pytest
-import torch
+
+try:
+    import mindspeed.megatron_adaptor  # noqa: F401 isort: skip  # must precede mbridge on NPU
+except ImportError:
+    pass
 
 from areal.api.alloc_mode import ModelAllocation
+from areal.infra.platforms import current_platform
 from areal.utils.network import find_free_ports
 from areal.utils.testing_utils import DENSE_MODEL_PATHS, MOE_MODEL_PATHS
 
@@ -29,7 +34,8 @@
     / "run_megatron_engine_vlm_distributed.py"
 ).resolve()
 
-CUDA_AVAILABLE = torch.cuda.is_available()
+# Detect any accelerator (CUDA or NPU)
+ACCELERATOR_AVAILABLE = current_platform.device_type in ("cuda", "npu")
 
 
 def _run_vlm_test(
@@ -122,7 +128,7 @@ def _run_vlm_test(
         {"VLM_MODEL_PATH": MOE_MODEL_PATHS["qwen3_vl_moe"]},
         id="qwen3_vl_moe",
         marks=pytest.mark.skipif(
-            torch.cuda.device_count() < 8,
+            current_platform.device_count() < 8,
             reason="Qwen3-VL-MoE-30B-A3B requires at least 8 GPUs",
         ),
     ),
@@ -131,7 +137,7 @@ def _run_vlm_test(
 
 @pytest.mark.gpu
 @pytest.mark.slow
-@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
+@pytest.mark.skipif(not ACCELERATOR_AVAILABLE, reason="No accelerator available")
 @pytest.mark.parametrize("model_env", _VLM_MODELS)
 def test_engine_initializes(model_env, tmp_path_factory):
     """Verify VLM engine detects vision model and loads processor."""
@@ -141,7 +147,7 @@ def test_engine_initializes(model_env, tmp_path_factory):
 
 @pytest.mark.gpu
 @pytest.mark.slow
-@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
+@pytest.mark.skipif(not ACCELERATOR_AVAILABLE, reason="No accelerator available")
 @pytest.mark.parametrize("model_env", _VLM_MODELS)
 def test_simple_forward(model_env, tmp_path_factory):
     """Verify forward pass with VLM inputs completes."""
@@ -151,7 +157,7 @@ def test_simple_forward(model_env, tmp_path_factory):
 
 @pytest.mark.gpu
 @pytest.mark.slow
-@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
+@pytest.mark.skipif(not ACCELERATOR_AVAILABLE, reason="No accelerator available")
 @pytest.mark.parametrize("model_env", _VLM_MODELS)
 def test_hf_save_load_weights(model_env, tmp_path_factory):
     """Verify save/load preserves VLM weights and saves processor."""
@@ -168,11 +174,11 @@ def test_hf_save_load_weights(model_env, tmp_path_factory):
 @pytest.mark.gpu
 @pytest.mark.multi_gpu
 @pytest.mark.slow
-@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
+@pytest.mark.skipif(not ACCELERATOR_AVAILABLE, reason="No accelerator available")
 @pytest.mark.parametrize("model_env", _VLM_MODELS)
 def test_train_tensor_parallel(model_env, tmp_path_factory):
     """VLM training with TP=2 to avoid single-device OOM."""
-    if torch.cuda.device_count() < 2:
+    if current_platform.device_count() < 2:
         pytest.skip("VLM TP training requires at least 2 GPUs")
     output = str(tmp_path_factory.mktemp("vlm_test") / "train_tp2.out")
     _run_vlm_test(
@@ -191,7 +197,7 @@ def test_train_tensor_parallel(model_env, tmp_path_factory):
 
 @pytest.mark.multi_gpu
 @pytest.mark.slow
-@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
+@pytest.mark.skipif(not ACCELERATOR_AVAILABLE, reason="No accelerator available")
 def test_qwen3vl_moe_expert_parallel(tmp_path_factory):
     """Forward smoke test for Qwen3-VL-MoE under ``(attn:d2t4|ffn:d2e4)``.
 
@@ -200,7 +206,7 @@ def test_qwen3vl_moe_expert_parallel(tmp_path_factory):
     ``InvalidAllocationModeError`` (attn=4 vs ffn=8). Validates the hybrid
     attn/ffn parser, EP-aware weight init, and VLM forward path.
     """
-    if torch.cuda.device_count() < 8:
+    if current_platform.device_count() < 8:
         pytest.skip("Qwen3-VL-MoE expert parallel requires 8 GPUs to run")
     output = str(
         tmp_path_factory.mktemp("test_output") / "qwen3vl_moe_expert_parallel.out"
@@ -215,15 +221,15 @@ def test_qwen3vl_moe_expert_parallel(tmp_path_factory):
 
 @pytest.mark.multi_gpu
 @pytest.mark.slow
-@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
+@pytest.mark.skipif(not ACCELERATOR_AVAILABLE, reason="No accelerator available")
 def test_qwen3vl_moe_dcp_save_load(tmp_path_factory):
     """DCP save/load round-trip for Qwen3-VL-MoE under ``(attn:d2p1t4|ffn:d1p1t2e4)``.
 
     Allocation: attn DP=2 PP=1 TP=4 (8 GPUs); ffn DP=1 PP=1 TP=2 EP=4 (8 GPUs).
     Drops the ``cp=2`` segment from the dense Qwen3-MoE analog because VLM
     forbids CP>1 (megatron_engine.py:347).
     """
-    if torch.cuda.device_count() < 8:
+    if current_platform.device_count() < 8:
         pytest.skip("Qwen3-VL-MoE DCP save load requires 8 GPUs to run")
     output = str(tmp_path_factory.mktemp("test_output") / "qwen3vl_moe_save_load.out")
     _run_vlm_test(

tests/torchrun/run_megatron_engine_vlm_distributed.py

@@ -19,7 +19,12 @@
 from pathlib import Path
 from typing import Any
 
-import torch
+try:
+    import mindspeed.megatron_adaptor  # noqa: F401  # must precede mbridge on NPU
+except ImportError:
+    pass
+
+import torch  # noqa: I001
 import torch.distributed as dist
 
 from areal.api import FinetuneSpec, SaveLoadMeta
@@ -31,6 +36,7 @@
     TrainEngineConfig,
 )
 from areal.engine import MegatronEngine
+from areal.infra.platforms import current_platform
 from areal.utils.data import broadcast_tensor_container
 from areal.utils.testing_utils import DENSE_MODEL_PATHS
 
@@ -135,7 +141,7 @@ def _make_input(engine: MegatronEngine) -> dict[str, Any]:
 
 
 def _cleanup(engine: MegatronEngine):
-    torch.cuda.synchronize()
+    current_platform.synchronize()
     dist.barrier()
     engine.destroy()