Merge branch 'hf-diverge-fix' into 'main'

Mark weights from vision encoder to be non-tensor-parallelizable to ensure gradients are correctly all-reduced

See merge request ADLR/megatron-lm!3190

Author: ko3n1g

megatron/core/distributed/finalize_model_grads.py

@@ -189,45 +189,6 @@ def _allreduce_embedding_grads(model: List[torch.nn.Module], config: Transformer
     _allreduce_position_embedding_grads(model, config)
 
 
-def _allreduce_layernorm_grads(model: List[torch.nn.Module], config: TransformerConfig):
-    """
-    All-reduce layernorm grads (for sequence parallelism).
-    """
-
-    # All-reduce layernorm parameters across model parallel nodes
-    # when sequence parallelism is used
-    if parallel_state.get_tensor_model_parallel_world_size() > 1 and (
-        config.sequence_parallel or config.qk_layernorm
-    ):
-        params = []
-        grads = []
-        for model_chunk in model:
-            ddp_config = model_chunk.ddp_config
-            for name, param in get_attr_wrapped_model(model_chunk, 'named_parameters')():
-                if param.requires_grad and (
-                    getattr(param, 'sequence_parallel', False)
-                    or 'q_layernorm' in name
-                    or 'k_layernorm' in name
-                ):
-                    params.append(param)
-                    grad_attr = _get_main_grad_attr(param, ddp_config.use_custom_fsdp)
-                    grad = getattr(param, grad_attr)
-                    grad = _unshard_if_dtensor(grad)
-                    grads.append(grad.data)
-        if grads:
-            coalesced = _flatten_dense_tensors(grads)
-            torch.distributed.all_reduce(
-                coalesced, group=parallel_state.get_tensor_model_parallel_group()
-            )
-            for param, buf, synced in zip(
-                params, grads, _unflatten_dense_tensors(coalesced, grads)
-            ):
-                buf.copy_(synced)
-                grad_attr = _get_main_grad_attr(param, ddp_config.use_custom_fsdp)
-                orig_grad = getattr(param, grad_attr)
-                setattr(param, grad_attr, _reshard_if_dtensor(buf, orig_grad))
-
-
 def _update_router_expert_bias(model: List[torch.nn.Module], config: TransformerConfig):
     """
     Update the expert bias of the router for a global batch.
@@ -256,6 +217,70 @@ def _update_router_expert_bias(model: List[torch.nn.Module], config: Transformer
         expert_bias.copy_(updated_expert_bias)
 
 
+def _allreduce_non_tensor_model_parallel_grads(
+    model: List[torch.nn.Module], config: TransformerConfig
+):
+    """
+    All-reduce both layernorm grads (for sequence parallelism) and
+    gradients from modules with average_gradients_across_tp_domain=True
+    across tensor-model-parallel ranks.
+    """
+    if parallel_state.get_tensor_model_parallel_world_size() <= 1:
+        return
+
+    params_sum = []
+    grads_sum = []
+    params_avg = []
+    grads_avg = []
+
+    for model_chunk in model:
+        ddp_config = model_chunk.ddp_config
+        for name, param in get_attr_wrapped_model(model_chunk, 'named_parameters')():
+            if param.requires_grad:
+                # Check if this param needs average reduction (average_gradients_across_tp_domain)
+                if getattr(param, "average_gradients_across_tp_domain", False):
+                    params_avg.append(param)
+                    grad_attr = _get_main_grad_attr(param, ddp_config.use_custom_fsdp)
+                    grad = getattr(param, grad_attr)
+                    grad = _unshard_if_dtensor(grad)
+                    grads_avg.append(grad.data)
+                # Check if this param needs sum reduction (sequence parallel or qk_layernorm)
+                elif (config.sequence_parallel and getattr(param, "sequence_parallel", False)) or (
+                    config.qk_layernorm and ("q_layernorm" in name or "k_layernorm" in name)
+                ):
+                    params_sum.append(param)
+                    grad_attr = _get_main_grad_attr(param, ddp_config.use_custom_fsdp)
+                    grad = getattr(param, grad_attr)
+                    grad = _unshard_if_dtensor(grad)
+                    grads_sum.append(grad.data)
+
+    # Loop grads and perform correct all-reduce
+    for params, grads, all_reduce_op in zip(
+        [params_sum, params_avg],
+        [grads_sum, grads_avg],
+        [torch.distributed.ReduceOp.SUM, torch.distributed.ReduceOp.AVG],
+    ):
+        if grads:
+            coalesced = _flatten_dense_tensors(grads)
+            torch.distributed.all_reduce(
+                coalesced, op=all_reduce_op, group=parallel_state.get_tensor_model_parallel_group()
+            )
+            for param, buf, synced in zip(
+                params, grads, _unflatten_dense_tensors(coalesced, grads)
+            ):
+                buf.copy_(synced)
+                grad_attr = _get_main_grad_attr(param, ddp_config.use_custom_fsdp)
+                orig_grad = getattr(param, grad_attr)
+                setattr(param, grad_attr, _reshard_if_dtensor(buf, orig_grad))
+
+
+"""
+This is an alias to _allreduce_non_tensor_model_parallel_grads that we must
+maintain for legacy tests. We can remove this proxy in mcore 0.14.
+"""
+_allreduce_layernorm_grads = _allreduce_non_tensor_model_parallel_grads
+
+
 def finalize_model_grads(model: List[torch.nn.Module], num_tokens: Optional[torch.Tensor] = None):
     """
     All-reduce all model grads across DP replicas, layernorm grads for sequence parallelism,
@@ -282,14 +307,14 @@ def finalize_model_grads(model: List[torch.nn.Module], num_tokens: Optional[torc
     if config.timers is not None:
         config.timers('conditional-embedder-grads-all-reduce').stop()
 
-    # All-reduce layer-norm grads (for sequence parallelism).
+    # All-reduce layer-norm grads (for sequence parallelism) and non-tensor parallel modules.
     if config.timers is not None:
-        config.timers('layernorm-grads-all-reduce', log_level=1).start(
+        config.timers('non-tensor-parallel-grads-all-reduce', log_level=1).start(
             barrier=config.barrier_with_L1_time
         )
-    _allreduce_layernorm_grads(model, config)
+    _allreduce_non_tensor_model_parallel_grads(model, config)
     if config.timers is not None:
-        config.timers('layernorm-grads-all-reduce').stop()
+        config.timers('non-tensor-parallel-grads-all-reduce').stop()
 
     # All-reduce embedding grads (for pipeline parallelism).
     if config.timers is not None:

megatron/core/models/huggingface/module.py

@@ -1,5 +1,6 @@
 # Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
 
+import torch
 from transformers import AutoConfig, AutoModel
 
 from megatron.core.transformer.module import MegatronModule
@@ -17,6 +18,20 @@ def set_input_tensor(self, input_tensor):
         """Dummy function for set_input_tensor"""
         self.input_tensor = input_tensor
 
+    def __setattr__(self, name: str, value):
+        """
+        Set average_gradients_across_tp_domain attribute true on all params so that during
+        finalize_model_grads an all-reduce is performed on this module’s gradients across
+        tensor parallel ranks. This keeps replicated weights synchronized and prevents drift
+        due to non determinism in HF models producing slightly different grads in replicated
+        models on the same inputs.
+        """
+        super().__setattr__(name, value)
+
+        if isinstance(value, torch.nn.Module):
+            for param in value.parameters(recurse=True):
+                setattr(param, "average_gradients_across_tp_domain", True)
+
 
 class AutoHuggingFaceModel(HuggingFaceModule):
     """

tests/unit_tests/distributed/test_finalize_model_grads.py

@@ -9,7 +9,7 @@
 from megatron.core import parallel_state
 from megatron.core.distributed import DistributedDataParallelConfig
 from megatron.core.distributed.finalize_model_grads import (
-    _allreduce_layernorm_grads,
+    _allreduce_non_tensor_model_parallel_grads,
     _allreduce_word_embedding_grads,
 )
 from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
@@ -67,7 +67,7 @@ def test_allreduce_layernorm_grads(self, freeze_model, tp_size):
             else:
                 param.grad = torch.ones_like(param)
 
-        _allreduce_layernorm_grads([self.model], self.transformer_config)
+        _allreduce_non_tensor_model_parallel_grads([self.model], self.transformer_config)
 
     @pytest.mark.parametrize(
         ("freeze_model", "pp_size", "share_embeddings"),