[li-lizhe]

What

GatedDeltaNet gathers the sequence before in_proj and turns sequence parallelism off on the projections 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, and Transformer Engine marks the parameters it wants reduced with the
same flag (layer_norm_weight of a LayerNormLinear). Only the module attribute was toggled, so on every
tensor parallel rank the gamma gradient was complete and still summed: with TP4 the gradient of
self_attention.in_proj.layer_norm_weight came out 4x the unsharded reference and reached the optimizer and
the global gradient clipping that way.

_set_linear_sequence_parallel now toggles the flag on the submodules and on their parameters, and
_restore_linear_sequence_parallel restores the modules while leaving the parameters at False. The
parameters deliberately stay off: the finalizer reads them after the backward, long after forward restored
the modules, and since the projections run without sequence parallelism every rank already holds the complete
gradient for them. align_corners-style behaviour of the projections themselves is untouched — the module
attribute is saved and restored exactly as before.

Fixes #241.

Tests

tests/test_gdn_linear_sequence_parallel_params.py (CPU, no GPU, no model weights) stands in for a
Transformer Engine LayerNormLinear (a module and a layer_norm_weight parameter that both carry the flag,
plus an unmarked weight) and calls the two helpers on it:

• the toggle marks the parameter and records it, the restore brings the module back and leaves the parameter
  off;
• a finalizer stub that applies Megatron's rule (param.main_grad * tp_size for a marked parameter) leaves the
  gamma gradient unscaled, which is the behaviour the issue asks for;
• a parameter without the flag, and a module without it, are left alone.

revision	result
main (23d95fd) + this test file	2 failed, 1 passed — assert True is False on the parameter marker, and the finalizer scales the gamma gradient by 4
this branch	3 passed

Verification

Run on a Linux CPU host, Python 3.14.6, torch 2.15.0.dev20260917+cpu, pytest 9.1.1, megatron-core from
PyPI, with the transformer_engine, triton and datasets imports stubbed (this machine has no GPU and no
Transformer Engine build) and mcore_bridge loaded as namespace packages so the package __init__ patch
chain is skipped. The helper functions under test are the real ones, imported from the patched file.

Checked before opening: no other open PR references #241 (/search/issues?q=repo:modelscope/mcore-bridge+is:pr+241
and the issue timeline).

Not tested: the TP4/CP2 reproduction from the issue and the full Qwen3.8-27B training runs (no GPU here),
and the repository's own test suite beyond this file (it needs the full Megatron + Transformer Engine
environment). The multiplier is a property of the flag the finalizer reads, which the test asserts directly.

[hjh0119]

Can this be replaced with the Megatron API?

[li-lizhe]

Thanks for the review — good point, and it is done in 9588175b.

The test no longer re-implements Megatron's rule. It now calls the real _allreduce_non_tensor_model_parallel_grads (megatron.core.distributed.finalize_model_grads) on two gloo ranks and reads the reduction off main_grad, so the assertion is Megatron's own code path rather than a stub of it. The Transformer Engine shaped LayerNormLinear stand-in stays, since the point of the test is the two sequence_parallel flags the toggle has to keep in sync.

To keep the pin meaningful I kept the pre-fix state as a control case: param-flag-off is what the fixed forward leaves behind (gamma stays at 1x, which is the issue's ask) and param-flag-on-control is the old state (tp_size times). If the parameter half of the toggle ever goes missing, the control pair stops differing and the test fails.

Verification (local, CPU, no GPU): megatron-core 0.16.1 on Python 3.11 with torch 2.7.1+cpu — pytest tests/test_gdn_linear_sequence_parallel_params.py -> 2 passed, 2 skipped. The two that pass are the finalizer cases; the two that skip are the toggle pins, which need the full Transformer Engine that mcore_bridge imports at module scope and which is not installable here. flake8 (repo config), isort and yapf (repo setup.cfg) are clean on the file.

One caveat to be explicit about: this is a CPU/gloo run of the reduction, not the real NCCL TP4 training job from the issue, so it pins the flag semantics and the reduction, not the end-to-end gamma magnitude on a GPU cluster.