fix(gdn): toggle sequence parallelism on the projection parameters too

src/mcore_bridge/model/modules/gated_delta_net.py

@@ -13,7 +13,7 @@
 from megatron.core.transformer.spec_utils import build_module
 from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint, sharded_state_dict_default
 from megatron.core.utils import deprecate_inference_params, nvtx_range_pop, nvtx_range_push
-from typing import List, Optional
+from typing import List, Optional, Union
 
 from mcore_bridge.config import ModelConfig
 
@@ -162,8 +162,18 @@ def _resolve_cu_seqlens(cu_seqlens_padded, cu_seqlens_actual, total_seq_len, cp_
             return None
         return cu_seqlens
 
-    def _set_linear_sequence_parallel(self, enabled: bool) -> list[tuple[nn.Module, bool]]:
-        saved: list[tuple[nn.Module, bool]] = []
+    def _set_linear_sequence_parallel(self, enabled: bool) -> list[tuple[Union[nn.Module, nn.Parameter], bool]]:
+        """Set ``sequence_parallel`` on the projections and on their parameters.
+
+        Transformer Engine reads the module attribute to choose between reduce-scatter and all-reduce for
+        the projection gradient, while Megatron's gradient finalizer reads ``param.sequence_parallel``
+        (``finalize_model_grads``) to decide whether that gradient still needs to be summed over the tensor
+        parallel group. Transformer Engine marks the parameters it wants reduced with the same flag, for
+        example ``layer_norm_weight`` in a ``LayerNormLinear``, so the two have to be toggled together:
+        with only the module attribute off, the projection all-reduces the gradient itself and the
+        finalizer sums the already complete copy again, which multiplies it by the tensor parallel size.
+        """
+        saved: list[tuple[Union[nn.Module, nn.Parameter], bool]] = []
         for name in ('in_proj', 'in_proj_qkvz', 'in_proj_ba', 'out_proj'):
             module = getattr(self, name, None)
             if module is None:
@@ -172,11 +182,23 @@ def _set_linear_sequence_parallel(self, enabled: bool) -> list[tuple[nn.Module,
                 if hasattr(sub, 'sequence_parallel'):
                     saved.append((sub, sub.sequence_parallel))
                     sub.sequence_parallel = enabled
+                for param in sub.parameters(recurse=False):
+                    if hasattr(param, 'sequence_parallel'):
+                        saved.append((param, param.sequence_parallel))
+                        param.sequence_parallel = enabled
         return saved
 
-    def _restore_linear_sequence_parallel(self, saved: list[tuple[nn.Module, bool]]) -> None:
+    def _restore_linear_sequence_parallel(self, saved: list[tuple[Union[nn.Module, nn.Parameter], bool]]) -> None:
+        """Restore the module attributes; leave the parameters at ``enabled=False``.
+
+        The parameters are deliberately not restored. ``finalize_model_grads`` reads them after the
+        backward, long after ``forward`` returned and restored the modules, and this layer gathers the
+        sequence before the projections, so every tensor parallel rank holds the complete gradient for the
+        parameters marked here and the finalizer must not sum them again.
+        """
         for module, enabled in saved:
-            module.sequence_parallel = enabled
+            if isinstance(module, nn.Module):
+                module.sequence_parallel = enabled
 
     def forward(
         self,

tests/test_gdn_linear_sequence_parallel_params.py

@@ -0,0 +1,95 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""The GDN sequence-parallel toggle has to reach the projection parameters.
+
+``GatedDeltaNet`` gathers the sequence before the projections and turns sequence parallelism off on them for
+that call, because Transformer Engine then all-reduces the projection gradient itself. Megatron's gradient
+finalizer decides from ``param.sequence_parallel`` (see ``finalize_model_grads``) whether a gradient still
+has to be summed over the tensor parallel group, so the flag has to be off on the parameters as well: with
+only the module attribute off, the finalizer sums the already complete gamma gradient once more and
+multiplies it by the tensor parallel size. Parameters that do not carry the flag are left alone.
+"""
+import pytest
+import torch
+import torch.nn as nn
+from types import SimpleNamespace
+
+
+class _FakeLayerNormLinear(nn.Module):
+    """Stand-in for a Transformer Engine ``LayerNormLinear``.
+
+    Transformer Engine flags the module and the parameters it wants reduced over the tensor parallel group
+    with ``sequence_parallel``; ``layer_norm_weight`` is the one the reported case loses in the finalizer.
+    """
+
+    def __init__(self, sequence_parallel: bool = True, mark_parameters: bool = True):
+        super().__init__()
+        self.sequence_parallel = sequence_parallel
+        self.layer_norm_weight = nn.Parameter(torch.ones(4))
+        self.weight = nn.Parameter(torch.ones(2, 4))
+        if mark_parameters:
+            self.layer_norm_weight.sequence_parallel = sequence_parallel
+
+
+def _gdn():
+    """Return ``GatedDeltaNet``, skipping the module when Megatron is not installed."""
+    pytest.importorskip('megatron.core')
+    from mcore_bridge.model.modules.gated_delta_net import GatedDeltaNet
+    return GatedDeltaNet
+
+
+def _projections(**modules):
+    """Only the four projection attributes the toggle walks are needed."""
+    attributes = {'in_proj': None, 'in_proj_qkvz': None, 'in_proj_ba': None, 'out_proj': None}
+    attributes.update(modules)
+    return SimpleNamespace(**attributes)
+
+
+def _finalize_layernorm_grad(param, tp_size):
+    """Megatron's branch: a sequence-parallel parameter has its complete gradient summed over the TP group."""
+    if getattr(param, 'sequence_parallel', False):
+        return param.main_grad * tp_size
+    return param.main_grad
+
+
+def test_toggle_marks_the_parameters_and_the_restore_keeps_them_off():
+    gated_delta_net = _gdn()
+    layer = _FakeLayerNormLinear()
+    model = _projections(in_proj=layer)
+    saved = gated_delta_net._set_linear_sequence_parallel(model, False)
+
+    assert layer.sequence_parallel is False
+    assert layer.layer_norm_weight.sequence_parallel is False
+    assert len(saved) == 2
+    assert saved[0][0] is layer
+    assert saved[1][0] is layer.layer_norm_weight
+
+    gated_delta_net._restore_linear_sequence_parallel(model, saved)
+    # The module goes back to its configured value, the parameter stays off: the finalizer reads it after
+    # the backward, long after the module attribute was restored.
+    assert layer.sequence_parallel is True
+    assert layer.layer_norm_weight.sequence_parallel is False
+
+
+def test_the_finalizer_leaves_the_gamma_gradient_unscaled():
+    gated_delta_net = _gdn()
+    layer = _FakeLayerNormLinear()
+    model = _projections(in_proj=layer)
+    saved = gated_delta_net._set_linear_sequence_parallel(model, False)
+    gated_delta_net._restore_linear_sequence_parallel(model, saved)
+
+    layer.layer_norm_weight.main_grad = torch.ones(4)
+    assert torch.equal(_finalize_layernorm_grad(layer.layer_norm_weight, 4), torch.ones(4))
+
+
+def test_parameters_without_the_flag_are_left_alone():
+    gated_delta_net = _gdn()
+    layer = _FakeLayerNormLinear(mark_parameters=False)
+    model = _projections(in_proj=layer)
+    saved = gated_delta_net._set_linear_sequence_parallel(model, False)
+
+    assert not hasattr(layer.weight, 'sequence_parallel')
+    assert len(saved) == 1
+    assert saved[0][0] is layer
+
+    gated_delta_net._restore_linear_sequence_parallel(model, saved)
+    assert layer.sequence_parallel is True