From 702e62f2b296678626e8288bfb81d81959f80b45 Mon Sep 17 00:00:00 2001 From: li-lizhe <82501526+li-lizhe@users.noreply.github.com> Date: Thu, 5 Nov 2026 04:28:49 +0800 Subject: [PATCH] [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 Co-authored-by: Aryan Kumar <156166681+aryan5v@users.noreply.github.com> --- .buildkite/scripts/unit_test.sh | 1 + fastvideo/distributed/communication_op.py | 13 ++- .../distributed/test_sp_warmup_device.py | 85 +++++++++++++++++++ 3 files changed, 96 insertions(+), 3 deletions(-) create mode 100644 fastvideo/tests/distributed/test_sp_warmup_device.py diff --git a/.buildkite/scripts/unit_test.sh b/.buildkite/scripts/unit_test.sh index 8d6a61bc1..0723cc97b 100644 --- a/.buildkite/scripts/unit_test.sh +++ b/.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 \ diff --git a/fastvideo/distributed/communication_op.py b/fastvideo/distributed/communication_op.py index 502284f06..eb6d207a7 100644 --- a/fastvideo/distributed/communication_op.py +++ b/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 diff --git a/fastvideo/tests/distributed/test_sp_warmup_device.py b/fastvideo/tests/distributed/test_sp_warmup_device.py new file mode 100644 index 000000000..1fa797da1 --- /dev/null +++ b/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 == []