@@ -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