diff --git a/cpp/tensorrt_llm/common/ncclUtils.cpp b/cpp/tensorrt_llm/common/ncclUtils.cpp index 5c82cb582d4d..4ced9615213f 100644 --- a/cpp/tensorrt_llm/common/ncclUtils.cpp +++ b/cpp/tensorrt_llm/common/ncclUtils.cpp @@ -24,6 +24,7 @@ #include #include #include +#include namespace { @@ -398,6 +399,7 @@ size_t NcclCommResourceManager::getResourceCount(ncclComm_t comm) const noexcept namespace { thread_local std::optional gGraphPoolOwner; +thread_local bool gWindowAllocationSuppressed = false; } // namespace NCCLWindowAllocator& NCCLWindowAllocator::getInstance() @@ -417,6 +419,11 @@ void NCCLWindowAllocator::setGraphPoolOwner(int64_t owner) gGraphPoolOwner = owner < 0 ? std::nullopt : std::optional{static_cast(owner)}; } +bool NCCLWindowAllocator::setWindowAllocationSuppressed(bool suppressed) +{ + return std::exchange(gWindowAllocationSuppressed, suppressed); +} + void NCCLWindowAllocator::releaseGraphPoolOwner(int64_t owner) { TLLM_CHECK_WITH_INFO(owner >= 0, "NCCL window graph-pool owner must be non-negative"); @@ -469,6 +476,12 @@ NCCLWindowBuffer NCCLWindowAllocator::requestBuffer(ncclComm_t comm, size_t size { return NCCLWindowBuffer(); } + if (gWindowAllocationSuppressed) + { + TLLM_LOG_TRACE("[NCCLUtil] NCCL window allocation suppressed on this thread for comm %p (requested: %zu)", + static_cast(comm), size); + return NCCLWindowBuffer(); + } TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL communicator cannot be null"); TLLM_CHECK_WITH_INFO(size > 0, "Buffer size must be greater than 0"); diff --git a/cpp/tensorrt_llm/common/ncclUtils.h b/cpp/tensorrt_llm/common/ncclUtils.h index cdd306369aa6..3d2416b4bd64 100644 --- a/cpp/tensorrt_llm/common/ncclUtils.h +++ b/cpp/tensorrt_llm/common/ncclUtils.h @@ -245,6 +245,11 @@ class NCCLWindowAllocator // still in use are returned when their deterministic tensor scope ends. void releaseGraphPoolOwner(int64_t owner); + // While set on this host thread, requestBuffer returns an invalid buffer without touching the + // pool, so callers fall back to plain memory. Used where ranks are not SPMD, such as PARALLEL + // autotuning. Returns the previous setting. + bool setWindowAllocationSuppressed(bool suppressed); + // Search for a buffer by pointer. Returns an invalid buffer if not found. // This matches the UBManager.search_buffer() interface. NCCLWindowBuffer searchBuffer(ncclComm_t comm, void* ptr) const; diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index ae96ce1b3e6a..1dcdcb153769 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -2498,6 +2498,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "end_nccl_window_tensor_scope.tracked(Tensor[] inputs, Tensor[] outputs, bool failed, int entry_depth) -> ()"); m.def("set_nccl_window_graph_owner(int owner) -> ()"); m.def("release_nccl_window_graph_owner(int owner) -> ()"); + m.def("set_nccl_window_allocation_suppressed(bool suppressed) -> bool"); m.def("get_nccl_window_buffer_release_epoch() -> int"); m.def("promote_nccl_window_buffer_releases(int device, int release_epoch) -> ()"); m.def( @@ -2558,6 +2559,12 @@ TORCH_LIBRARY_IMPL(trtllm, CatchAll, m) m.impl("release_nccl_window_graph_owner", [](int64_t owner) { tensorrt_llm::common::nccl_util::NCCLWindowAllocator::getInstance().releaseGraphPoolOwner(owner); }); + m.impl("set_nccl_window_allocation_suppressed", + [](bool suppressed) + { + return tensorrt_llm::common::nccl_util::NCCLWindowAllocator::getInstance().setWindowAllocationSuppressed( + suppressed); + }); m.impl("get_nccl_window_buffer_release_epoch", []() { @@ -2574,6 +2581,7 @@ TORCH_LIBRARY_IMPL(trtllm, CatchAll, m) #else m.impl("set_nccl_window_graph_owner", [](int64_t) {}); m.impl("release_nccl_window_graph_owner", [](int64_t) {}); + m.impl("set_nccl_window_allocation_suppressed", [](bool) { return false; }); m.impl("get_nccl_window_buffer_release_epoch", []() { return int64_t{0}; }); m.impl("promote_nccl_window_buffer_releases", [](int64_t, int64_t) {}); #endif diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index 8efca0be5a22..03c8931084b4 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -34,8 +34,8 @@ import tensorrt_llm from tensorrt_llm._torch.distributed import Distributed from tensorrt_llm._torch.nccl_window_graph import ( - abandon_nccl_window_graph_owner, nccl_window_graph_capture, - release_nccl_window_graph_owner) + abandon_nccl_window_graph_owner, nccl_window_allocation_suppressed, + nccl_window_graph_capture, release_nccl_window_graph_owner) from tensorrt_llm._torch.nccl_window_tensor_scope import \ nccl_window_tensor_scope from tensorrt_llm._utils import confidential_compute_enabled, nvtx_range @@ -1332,7 +1332,14 @@ def _profile_runners( tuning_config: TuningConfig, **kwargs, ) -> float: - with nccl_window_tensor_scope(input_tensors) as scope: + # PARALLEL ranks profile different tactics, so window requests here would diverge the + # per-rank pools and let a later request register a window collectively on some ranks only. + window_guard = (nccl_window_allocation_suppressed() + if tuning_config.distributed_tuning_strategy + == DistributedTuningStrategy.PARALLEL + and self._is_distributed() else + contextlib.nullcontext()) + with window_guard, nccl_window_tensor_scope(input_tensors) as scope: scope.escape(input_tensors) if tuning_config.inputs_pre_hook is not None: input_tensors = tuning_config.inputs_pre_hook(input_tensors) diff --git a/tensorrt_llm/_torch/nccl_window_graph.py b/tensorrt_llm/_torch/nccl_window_graph.py index 84d4fb706bb2..3a0709ca89a1 100644 --- a/tensorrt_llm/_torch/nccl_window_graph.py +++ b/tensorrt_llm/_torch/nccl_window_graph.py @@ -106,3 +106,18 @@ def release_nccl_window_graph_owner(pool: Any) -> None: if owner is not None: torch.ops.trtllm.release_nccl_window_graph_owner(owner) del _pool_owners[key] + + +@contextlib.contextmanager +def nccl_window_allocation_suppressed() -> Iterator[None]: + """Serve NCCL-window requests on this thread from plain memory. + + The window pool is left untouched, so work that does not run in SPMD order + across ranks cannot make a later request miss on only some ranks and enter + the collective registration alone. + """ + previous = torch.ops.trtllm.set_nccl_window_allocation_suppressed(True) + try: + yield + finally: + torch.ops.trtllm.set_nccl_window_allocation_suppressed(previous) diff --git a/tests/unittest/_torch/multi_gpu/test_allocate_output_buffer_kinds.py b/tests/unittest/_torch/multi_gpu/test_allocate_output_buffer_kinds.py index 5eb80ee2aba5..9baecff37331 100644 --- a/tests/unittest/_torch/multi_gpu/test_allocate_output_buffer_kinds.py +++ b/tests/unittest/_torch/multi_gpu/test_allocate_output_buffer_kinds.py @@ -29,6 +29,13 @@ import tensorrt_llm import tensorrt_llm.bindings.internal.userbuffers as ub +from tensorrt_llm._torch.autotuner import ( + AutoTuner, + DistributedTuningStrategy, + TunableRunner, + TuningConfig, + autotune, +) from tensorrt_llm._torch.distributed import ( AllReduce, AllReduceFusionOp, @@ -222,6 +229,56 @@ def _allocate(ref, kind, tp_group): return True +def _run_parallel_tuning_then_window_allreduce(tp_size, tp_rank, seq_len, hidden_size): + """PARALLEL-tune a window-output runner, then all-reduce in a fresh window on every rank.""" + from tensorrt_llm.bindings.internal.thop import BufferKind + + def _allocate(ref, tp_group): + return torch.ops.trtllm.allocate_output(ref, int(BufferKind.NCCL_WINDOW), tp_group) + + class WindowOutputRunner(TunableRunner): + def __init__(self, tp_group): + self.tp_group = tp_group + + def get_valid_tactics(self, inputs, profile, **kwargs): + return [0] + + def forward(self, inputs, *, tactic=-1, **kwargs): + out, _ = _allocate(inputs[0], self.tp_group) + return out.copy_(inputs[0]) + + mapping = Mapping(world_size=tp_size, tp_size=tp_size, rank=tp_rank) + ref = torch.ones(seq_len, hidden_size, dtype=torch.bfloat16, device="cuda") + with discard_nccl_window_tensor_outputs(ref): + _, actual_kind_int = _allocate(ref, mapping.tp_group) + if actual_kind_int != int(BufferKind.NCCL_WINDOW): + return _SKIP + + tuner = AutoTuner.get() + tuner.clear_cache() + tuner.setup_distributed_state(mapping) + try: + # A single tactic is profiled by rank 0 only. + with autotune(): + tuner.choose_one( + "test_parallel_tuning_then_window_allreduce", + [WindowOutputRunner(mapping.tp_group)], + TuningConfig(distributed_tuning_strategy=DistributedTuningStrategy.PARALLEL), + [ref], + ) + finally: + tuner.setup_distributed_state(Mapping(), None) + + allreduce = AllReduce(mapping=mapping, strategy=AllReduceStrategy.NCCL_SYMMETRIC) + with discard_nccl_window_tensor_outputs(ref): + out, actual_kind_int = _allocate(ref, mapping.tp_group) + assert actual_kind_int == int(BufferKind.NCCL_WINDOW) + out.copy_(ref) + result = allreduce(out, all_reduce_params=AllReduceParams(fusion_op=AllReduceFusionOp.NONE)) + torch.testing.assert_close(result, torch.full_like(result, tp_size), atol=1e-2, rtol=1e-2) + return True + + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs") @pytest.mark.parametrize( "seq_len,hidden_size", @@ -297,3 +354,18 @@ def test_allreduce_nccl_window_buffer(seq_len, hidden_size, mpi_pool_executor): ) for r in results: assert r is True or r == _SKIP, f"Unexpected result: {r}" + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs") +@pytest.mark.timeout(300) +@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True) +def test_parallel_tuning_keeps_nccl_windows_symmetric(mpi_pool_executor): + tp_size = mpi_pool_executor.num_workers + results = list( + mpi_pool_executor.map( + run_single_rank, + *zip(*[(tp_size, _run_parallel_tuning_then_window_allreduce, 256, 4096)] * tp_size), + ) + ) + for r in results: + assert r is True or r == _SKIP, f"Unexpected result: {r}"