Skip to content

Commit 0ff7b4a

Browse files
authored
[https://nvbugs/6264844][fix] Fix wrong NCCL fallback in nemotron-h (#15294)
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
1 parent 6a4d8ca commit 0ff7b4a

1 file changed

Lines changed: 6 additions & 0 deletions

File tree

tensorrt_llm/_torch/models/modeling_nemotron_h.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -277,9 +277,11 @@ def _moe(name):
277277
)
278278

279279
if reduce_output:
280+
# AllReduce needs dtype at construction to build fused MNNVL paths.
280281
self.allreduce = AllReduce(
281282
mapping=model_config.mapping,
282283
strategy=model_config.allreduce_strategy,
284+
dtype=config.torch_dtype,
283285
)
284286
else:
285287
self.allreduce = None
@@ -488,9 +490,11 @@ def __init__(
488490
)
489491

490492
if fuse_allreduce_norm and layer_idx > 0:
493+
# AllReduce needs dtype at construction to build fused MNNVL paths.
491494
self.pre_allreduce = AllReduce(
492495
mapping=model_config.mapping,
493496
strategy=model_config.allreduce_strategy,
497+
dtype=config.torch_dtype,
494498
)
495499

496500
# Mixer creation. The fuse_allreduce_norm optimization is orthogonal
@@ -717,9 +721,11 @@ def __init__(self, model_config: NemotronHModelConfig):
717721

718722
# AllReduce for fusing with final norm (after last layer's mixer)
719723
if self.fuse_allreduce_norm:
724+
# AllReduce needs dtype at construction to build fused MNNVL paths.
720725
self.final_allreduce = AllReduce(
721726
mapping=model_config.mapping,
722727
strategy=model_config.allreduce_strategy,
728+
dtype=config.torch_dtype,
723729
)
724730

725731
def forward(

0 commit comments

Comments
 (0)