[bugfix] Use NPU current streams in PyHcclCommunicator

.buildkite/scripts/unit_test.sh

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

fastvideo/distributed/device_communicators/pyhccl.py

@@ -7,7 +7,6 @@
                                                                        hcclUniqueId)
 from fastvideo.distributed.utils import StatelessProcessGroup
 from fastvideo.logger import init_logger
-from fastvideo.utils import current_stream
 
 logger = init_logger(__name__)
 
@@ -99,7 +98,7 @@ def __init__(
         with torch.npu.device(device):
             self.comm: hcclComm_t = self.hccl.hcclCommInitRank(self.world_size, self.unique_id, self.rank)
 
-            stream = current_stream()
+            stream = torch.npu.current_stream(self.device)
             # A small all_reduce for warmup.
             data = torch.zeros(1, device=device)
             self.all_reduce(data)
@@ -118,7 +117,7 @@ def all_reduce(self, in_tensor: torch.Tensor, op: ReduceOp = ReduceOp.SUM, strea
         out_tensor = torch.empty_like(in_tensor)
 
         if stream is None:
-            stream = current_stream()
+            stream = torch.npu.current_stream(self.device)
         self.hccl.hcclAllReduce(buffer_type(in_tensor.data_ptr()), buffer_type(out_tensor.data_ptr()),
                                 in_tensor.numel(), hcclDataTypeEnum.from_torch(in_tensor.dtype),
                                 hcclRedOpTypeEnum.from_torch(op), self.comm, aclrtStream_t(stream.npu_stream))
@@ -130,7 +129,7 @@ def broadcast(self, tensor: torch.Tensor, src: int, stream=None):
         assert tensor.device == self.device, (f"this hccl communicator is created to work on {self.device}, "
                                               f"but the input tensor is on {tensor.device}")
         if stream is None:
-            stream = current_stream()
+            stream = torch.npu.current_stream(self.device)
         buffer = buffer_type(tensor.data_ptr())
         self.hccl.hcclBroadcast(buffer, tensor.numel(), hcclDataTypeEnum.from_torch(tensor.dtype), src, self.comm,
                                 aclrtStream_t(stream.npu_stream))

fastvideo/tests/distributed/test_pyhccl_stream.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+from unittest.mock import Mock, call
+
+import pytest
+import torch
+
+from fastvideo.distributed.device_communicators import pyhccl
+from fastvideo.distributed.device_communicators.pyhccl import PyHcclCommunicator
+from fastvideo.distributed.utils import StatelessProcessGroup
+
+
+@pytest.mark.parametrize("operation", ["all_reduce", "broadcast"])
+@pytest.mark.parametrize("explicit_stream", [False, True])
+def test_pyhccl_uses_communicator_device_current_stream(monkeypatch, operation, explicit_stream):
+    stream = SimpleNamespace(npu_stream=42)
+    current_stream = Mock(return_value=stream)
+    monkeypatch.setattr(torch, "npu", SimpleNamespace(current_stream=current_stream), raising=False)
+    communicator = object.__new__(PyHcclCommunicator)
+    communicator.disabled = False
+    communicator.device = torch.device("cpu")
+    communicator.comm = 1
+    communicator.hccl = SimpleNamespace(hcclAllReduce=Mock(), hcclBroadcast=Mock())
+    tensor = torch.ones(4)
+
+    kwargs = {"stream": stream} if explicit_stream else {}
+    if operation == "all_reduce":
+        communicator.all_reduce(tensor, **kwargs)
+        call = communicator.hccl.hcclAllReduce.call_args
+    else:
+        communicator.broadcast(tensor, src=0, **kwargs)
+        call = communicator.hccl.hcclBroadcast.call_args
+
+    if explicit_stream:
+        current_stream.assert_not_called()
+    else:
+        current_stream.assert_called_once_with(communicator.device)
+    assert call.args[-1].value == stream.npu_stream
+
+
+def test_pyhccl_warmup_uses_and_synchronizes_communicator_device_stream(monkeypatch):
+    stream = SimpleNamespace(npu_stream=42, synchronize=Mock())
+    current_stream = Mock(return_value=stream)
+    device_context = Mock(side_effect=lambda device: nullcontext())
+    monkeypatch.setattr(torch, "npu", SimpleNamespace(current_stream=current_stream, device=device_context), raising=False)
+    hccl = SimpleNamespace(hcclGetUniqueId=Mock(return_value=object()),
+                           hcclCommInitRank=Mock(return_value=1), hcclAllReduce=Mock())
+    monkeypatch.setattr(pyhccl, "HCCLLibrary", Mock(return_value=hccl))
+    group = StatelessProcessGroup(rank=0, world_size=2, store=Mock())
+    monkeypatch.setattr(group, "broadcast_obj", lambda value, src: value)
+    device = torch.device("cpu")
+
+    communicator = PyHcclCommunicator(group, device)
+
+    assert communicator.available
+    assert current_stream.call_args_list == [call(device), call(device)]
+    assert device_context.call_args_list == [call(device), call(device)]
+    hccl.hcclCommInitRank.assert_called_once_with(2, communicator.unique_id, 0)
+    hccl.hcclAllReduce.assert_called_once()
+    assert hccl.hcclAllReduce.call_args.args[-1].value == stream.npu_stream
+    stream.synchronize.assert_called_once_with()