[bugfix] SP warmup: use the rank's own device instead of hardcoded CUDA (#1988)

Signed-off-by: li-lizhe <147392333@qq.com>
Co-authored-by: aryan5v <aryan5v@users.noreply.github.com>
Co-authored-by: Aryan Kumar <156166681+aryan5v@users.noreply.github.com>

.buildkite/scripts/unit_test.sh

@@ -24,6 +24,7 @@ exec pytest \
   ./fastvideo/tests/stages/ \
   ./fastvideo/tests/ops/ \
   ./fastvideo/tests/worker/ \
+  ./fastvideo/tests/distributed/test_sp_warmup_device.py \
   ./fastvideo/tests/training/test_runner.py \
   ./fastvideo/tests/training/test_trackers.py \
   ./fastvideo/tests/training/test_ltx2_rope_fps.py \

fastvideo/distributed/communication_op.py

@@ -99,7 +99,8 @@ def warmup_sequence_parallel_communication(device: torch.device | None = None) -
     trigger the initialization upfront, before the first real forward pass.
     
     Args:
-        device: Device to use for warmup tensors. If None, uses CUDA device 0.
+        device: Device to use for warmup tensors. If None, uses the device type of the
+            accelerator the rank was placed on.
     """
     global _sp_warmup_done
 
@@ -114,8 +115,14 @@ def warmup_sequence_parallel_communication(device: torch.device | None = None) -
         _sp_warmup_done = True
         return
 
+    from fastvideo.platforms import current_platform
+    device_module = current_platform.get_torch_device()
+
     if device is None:
-        device = torch.device("cuda")
+        # The warmup tensors go through the SP collectives, so they have to live on the
+        # accelerator this rank runs on: a "cuda" default is not a valid device on a build
+        # without CUDA (Ascend NPU, MPS, ...) and the warmup would fail to allocate.
+        device = torch.device(current_platform.device_type)
 
     logger.info("Warming up sequence parallel communication (SP=%d)...", sp_world_size)
 
@@ -149,7 +156,7 @@ def warmup_sequence_parallel_communication(device: torch.device | None = None) -
     _ = sequence_model_parallel_all_gather(dummy3, dim=2)
 
     # Synchronize to ensure warmup completes
-    torch.cuda.synchronize(device)
+    device_module.synchronize(device)
 
     # Clean up
     del dummy, dummy2, dummy3

fastvideo/tests/distributed/test_sp_warmup_device.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import torch
+
+import fastvideo.platforms as platforms
+from fastvideo.distributed import communication_op
+
+
+class _RecordingSpGroup:
+    def __init__(self) -> None:
+        self.tensors: list[torch.Tensor] = []
+
+    def all_to_all_4D(self, tensor: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1) -> torch.Tensor:
+        self.tensors.append(tensor)
+        return tensor
+
+    def all_gather(self, tensor: torch.Tensor, dim: int = -1) -> torch.Tensor:
+        self.tensors.append(tensor)
+        return tensor
+
+
+class _RecordingDeviceModule:
+    def __init__(self) -> None:
+        self.synchronized: list = []
+
+    def synchronize(self, device=None) -> None:
+        self.synchronized.append(device)
+
+
+class _RecordingPlatform:
+    def __init__(self, device_type: str) -> None:
+        self.device_type = device_type
+        self.device_module = _RecordingDeviceModule()
+
+    def get_torch_device(self) -> _RecordingDeviceModule:
+        return self.device_module
+
+
+def _patch_warmup_inputs(monkeypatch, platform: _RecordingPlatform, group: _RecordingSpGroup) -> None:
+    # Set the lazy source, not the public attribute: a pinned current_platform
+    # would leak into later tests that replace _current_platform.
+    monkeypatch.setattr(platforms, "_current_platform", platform)
+    if "current_platform" in vars(platforms):
+        monkeypatch.delitem(vars(platforms), "current_platform")
+    monkeypatch.setattr(communication_op, "model_parallel_is_initialized", lambda: True)
+    monkeypatch.setattr(communication_op, "get_sp_world_size", lambda: 2)
+    monkeypatch.setattr(communication_op, "get_sp_group", lambda: group)
+    monkeypatch.setattr(communication_op, "_sp_warmup_done", False)
+
+
+def test_warmup_uses_the_active_platform_device(monkeypatch) -> None:
+    platform = _RecordingPlatform(device_type="cpu")
+    group = _RecordingSpGroup()
+    _patch_warmup_inputs(monkeypatch, platform, group)
+
+    communication_op.warmup_sequence_parallel_communication()
+
+    assert group.tensors, "the warmup must exercise the SP collectives"
+    assert {tensor.device.type for tensor in group.tensors} == {"cpu"}
+    assert platform.device_module.synchronized == [torch.device("cpu")]
+
+
+def test_warmup_keeps_an_explicit_device(monkeypatch) -> None:
+    platform = _RecordingPlatform(device_type="cuda")
+    group = _RecordingSpGroup()
+    _patch_warmup_inputs(monkeypatch, platform, group)
+
+    communication_op.warmup_sequence_parallel_communication(torch.device("cpu"))
+
+    assert {tensor.device.type for tensor in group.tensors} == {"cpu"}
+    assert platform.device_module.synchronized == [torch.device("cpu")]
+
+
+def test_warmup_is_skipped_without_sequence_parallelism(monkeypatch) -> None:
+    platform = _RecordingPlatform(device_type="cpu")
+    group = _RecordingSpGroup()
+    _patch_warmup_inputs(monkeypatch, platform, group)
+    monkeypatch.setattr(communication_op, "get_sp_world_size", lambda: 1)
+
+    communication_op.warmup_sequence_parallel_communication()
+
+    assert group.tensors == []
+    assert platform.device_module.synchronized == []