From 44ae5a496f13f13cab917027662a1bd27dc8af51 Mon Sep 17 00:00:00 2001 From: peihu-nv <259410613+peihu-nv@users.noreply.github.com> Date: Tue, 3 Mar 2026 11:06:59 -0800 Subject: [PATCH] Avoid NCCL_SYMMETRIC deadlock on OOM and reduce peak memory on DeepSeek MoE Signed-off-by: peihu-nv <259410613+peihu-nv@users.noreply.github.com> --- tensorrt_llm/_torch/custom_ops/torch_custom_ops.py | 6 +++++- tensorrt_llm/_torch/models/modeling_deepseekv3.py | 6 ++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index f88d13ddbd32..2b0f9b6b2c03 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -1769,7 +1769,11 @@ def forward( do_preparation=True) return input if tactic == -1: - tactic = AllReduceStrategy.NCCL_SYMMETRIC.value + # tactic == -1 means the autotuner cache missed for this shape, + # so we fall back to plain NCCL instead of NCCL_SYMMETRIC. + # NCCL_SYMMETRIC requires ncclMemAlloc which can fail asymmetrically + # across ranks under OOM, causing a deadlock at ncclCommWindowRegister. + tactic = AllReduceStrategy.NCCL.value return torch.ops.trtllm.allreduce( input, diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index 2318666ace0a..7c1a18173368 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -139,7 +139,8 @@ def weight_dequant(x: torch.Tensor, @torch.compile(dynamic=True) def moe_reduce_add_shared_output(routed_output, shared_output): routed_output = torch.sum(routed_output, dim=1, keepdim=False) - return shared_output + routed_output + # In-place add to avoid allocating a temporary tensor, reducing peak memory + return shared_output.add_(routed_output) class DeepseekV3WeightLoader: @@ -1157,7 +1158,8 @@ def _compute_routed_output(): else: assert shared_output.size() == routed_output.size( ), 'unmatched tensor shape' - final_hidden_states = shared_output + routed_output + # In-place add to avoid allocating a temporary tensor, reducing peak memory + final_hidden_states = shared_output.add_(routed_output) if not self.use_dp and self.mapping.tp_size > 1: final_hidden_states = self.allreduce(