From 2b11941650f4b4891002491ce7dbaaecdeb4b149 Mon Sep 17 00:00:00 2001 From: wenchenvincent <32376000+wenchenvincent@users.noreply.github.com> Date: Tue, 25 Aug 2026 02:41:35 +0000 Subject: [PATCH 1/3] Do not mark replicated weights as tensor-model-parallel Linear.reset_parameters marked every weight is_parallel=True regardless of parallel_mode. For parallel_mode=None the weight is replicated on every TP rank, so downstream consumers that use the attribute to de-duplicate -- notably Megatron's param_is_not_tensor_parallel_duplicate(), which gates get_grads_for_norm() -- admit it to the global gradient norm once per rank instead of once. The norm is assembled as a sum of squares across ranks, so the contribution is added tp_size times. Inflation follows sqrt(1 + (TP-1)*f), where f is the replicated weights share of the true squared norm. Measured at TP=8 on MI355X: DeepSeek-V4-Flash 3.763 -> 2.472 (1.52x, f~0.19) DeepSeek-V3 12.228 -> 9.011 (1.36x, f~0.12) This is primarily a diagnostic bug: the reported norm is what practitioners use to judge training health, tune --clip-grad and compare against reference curves. The effect on weights is optimizer-dependent -- clipping is a global uniform rescale, which Adam largely absorbs and Muon absorbs exactly, while SGD sees it in full. Affects architectures that deliberately replicate weights carrying real gradient energy. Megatron MLA passes parallel_mode=duplicated for q_down_proj/kv_down_proj, so DeepSeek-V2/V3/V3.2-family models take this path. Restores TE own default: _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS declares tensor_model_parallel=False for a non-parallel tensor. Introduced upstream in 044903374 (2023-02-10). --- transformer_engine/pytorch/module/linear.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index c4c9318b7b..84df0f4a8b 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1843,7 +1843,12 @@ def reset_parameters(self, defer_init=False): for weight in self.weight_names: set_tensor_model_parallel_attributes( tensor=getattr(self, weight), - is_parallel=True, + # A weight is only tensor-model-parallel when the layer is. For + # parallel_mode=None the weight is replicated on every TP rank, and + # marking it parallel makes downstream consumers (e.g. Megatron's + # param_is_not_tensor_parallel_duplicate) admit it to the global + # gradient norm once per rank instead of once. + is_parallel=self.parallel_mode is not None, dim=1 if self.parallel_mode == "row" else 0, stride=1, ) From cd50a1c5630d754aba9db71e35b79e63833d0efa Mon Sep 17 00:00:00 2001 From: wenchenvincent <32376000+wenchenvincent@users.noreply.github.com> Date: Tue, 25 Aug 2026 02:41:46 +0000 Subject: [PATCH 2/3] docs: explain the replicated-weight gradient-norm overcount Mechanism, the sqrt(1 + (TP-1)*f) scaling law, which axis is affected, optimizer-dependent impact (SGD full / Adam second-order / Muon none), affected architectures, history bisect, fix and caveats. Evidence from two MLA models at TP=8 on MI355X: DeepSeek-V4-Flash (1.52x) and DeepSeek-V3 (1.36x). The DeepSeek-V3 run uses Megatron pretrain_gpt.py with --multi-latent-attention --transformer-impl transformer_engine and no third-party model code. --- docs/replicated_weight_gradnorm.md | 281 +++++++++++++++++++++++++++++ 1 file changed, 281 insertions(+) create mode 100644 docs/replicated_weight_gradnorm.md diff --git a/docs/replicated_weight_gradnorm.md b/docs/replicated_weight_gradnorm.md new file mode 100644 index 0000000000..5e09539860 --- /dev/null +++ b/docs/replicated_weight_gradnorm.md @@ -0,0 +1,281 @@ +# Replicated weights are counted `tp_size` times in the global gradient norm + +## Summary + +`transformer_engine.pytorch.Linear` marks **every** weight as tensor-model-parallel, +including weights that are explicitly replicated (`parallel_mode=None`, which is what +Megatron's `"duplicated"` maps to). Megatron reads that attribute to decide whether a +parameter enters the global gradient-norm sum once, or once per TP rank. Replicated +weights are therefore counted `tp_size` times and the reported gradient norm is +inflated. + +Measured on two MLA models at TP=8, each with a constant per-iteration ratio: + +| model | stock | fixed | inflation | +|---|---|---|---| +| DeepSeek-V4-Flash (4 layers) | 3.763 | 2.472 | **1.52×** | +| DeepSeek-V3 (6 layers, Megatron's own MLA) | 12.228 | 9.011 | **1.36×** | + +The DeepSeek-V3 run uses Megatron's `pretrain_gpt.py` with +`--multi-latent-attention --transformer-impl transformer_engine` and no third-party +model code, so the behaviour is reproducible with upstream components alone. + +**This is primarily a diagnostic bug.** The gradient norm is the number practitioners +use to judge training health, tune `--clip-grad`, detect instability, and compare +against reference curves — and it is wrong by 52% here. Its effect on the *weights* +is much smaller, and depends entirely on the optimizer (see +[Impact on training](#impact-on-training)). + +## Mechanism + +`transformer_engine/pytorch/module/linear.py`, `Linear.reset_parameters`: + +```python +for weight in self.weight_names: + set_tensor_model_parallel_attributes( + tensor=getattr(self, weight), + is_parallel=True, # unconditional + dim=1 if self.parallel_mode == "row" else 0, + stride=1, + ) +``` + +`parallel_mode` selects the *dim* but never whether to mark at all, so a non-parallel +linear is still labelled `tensor_model_parallel=True`. + +Megatron consumes it in `param_is_not_tensor_parallel_duplicate()`: + +```python +if hasattr(param, "tensor_model_parallel") and param.tensor_model_parallel: + return True # -> admitted on EVERY TP rank +return tp_group.rank() == 0 # -> replicated: admitted once +``` + +which gates `get_grads_for_norm()`. The norm is then assembled as a sum of squares +across ranks, so a replicated weight's contribution is added `tp_size` times: + +```python +total_norm = local_l2_norm(grads_for_norm) ** 2 +all_reduce(total_norm, op=SUM, group=grad_stats_parallel_group) +total_norm = total_norm ** 0.5 +``` + +Two details confirm the intent: + +* TE never *reads* the attribute back — it is set purely for downstream consumers. +* TE's own `_MODEL_PARALLEL_ATTRIBUTE_DEFAULTS` declares + `{'tensor_model_parallel': False, 'partition_dim': -1, 'partition_stride': 1}`; + the code overrides its own stated default for the non-parallel case. + +Megatron does not correct it afterwards: the post-init loop in +`megatron/core/extensions/transformer_engine.py` sets only `allreduce` and +`sequence_parallel`. + +Note there is no separate "duplicated" state. The attribute set is entirely TP-scoped +and TE's `Linear` has no data-parallel concept at all (`sequence_parallel`, `tp_group`, +`tp_size`, `parallel_mode`). Replicated *is* `tensor_model_parallel=False`. + +## How large is it + +$$\text{inflation} = \sqrt{1 + (\text{TP} - 1)\,f}$$ + +where `f` is the duplicated weights' share of the **true squared norm** — gradient +*energy*, not parameter count. The distinction matters: in DeepSeek-V4-Flash the +replicated projections are ~1.5% of parameters but hold ~19% of gradient energy, +because they see the full hidden state and every head's gradient flows back through +them. + +| f | TP=2 | TP=4 | TP=8 | TP=16 | TP=32 | +|---|---|---|---|---|---| +| 1.5% (typical replicated weight) | 1.01 | 1.02 | 1.05 | 1.09 | 1.12 | +| 19% (MLA down-projections) | 1.09 | 1.25 | **1.53** | 1.96 | 2.62 | + +Inverting the measured ratios gives `f` directly, and both models land far above their +parameter share (~0.3%), confirming the effect tracks gradient energy rather than +parameter count: + +| model | measured ratio @ TP=8 | implied f | +|---|---|---| +| DeepSeek-V4-Flash | 1.522 | 18.8% | +| DeepSeek-V3 | 1.361 | 12.2% | + +So the same bug is invisible in most models and significant in MLA at large TP. + +## Which axis is affected + +Only tensor parallelism. + +| axis | protection | +|---|---| +| DP | Sharding. With `--use-distributed-optimizer` each rank's `get_parameters()` returns only its shard, so a gradient enters the sum on exactly one DP rank. Without it, the norm group is the model-parallel group, excluding DP. | +| TP | `param_is_not_tensor_parallel_duplicate()` — the attribute above, and the **only** guard on this axis. | + +## Impact on training + +Megatron's `--clip-grad` is *norm*-based and *global*: one scalar for the whole model, +applied uniformly. + +```python +clip_coeff = max_norm / (total_norm + 1.0e-6) +if clip_coeff < 1.0: + multi_tensor_scale(grads, grads, clip_coeff) +``` + +This is a **rescale, not a clip** — direction and the relative weighting between +parameters are preserved exactly. Three properties together bound the damage: + +1. **global**, not per-parameter — a wrong scalar is a magnitude error, not a + relative-weighting error; +2. **norm-based**, not value-based — no element-wise truncation, so direction is + untouched; +3. paired with a **scale-invariant optimizer**, the magnitude error is absorbed. + +It also only engages at all when the inflated norm crosses the threshold: + +| condition | effect | +|---|---| +| `--clip-grad 0` | none — the norm is not even computed | +| both true and inflated norm below threshold | none — `clip_coeff >= 1`, gated off | +| inflated norm crosses the threshold | gradients scaled by `1/inflation` | + +DeepSeek-V4-Flash is in the third regime: true norm ~2.47 against `--clip-grad 1.0`, +so clipping fires every step, with coefficient `0.266` instead of `0.405` — gradients +34% smaller than intended. + +What that does to the weights depends on the optimizer: + +| optimizer | scale invariance | impact | +|---|---|---| +| **SGD** | none — update is linear in the gradient | **full**: 34% smaller steps. Also, weight decay is coupled and applied *after* clipping, so the decay term is not scaled and becomes 1.52× stronger relative to the gradient. | +| **Adam / AdamW** | approximate | **second-order**. `m -> c·m`, `v -> c²·v`, so `update = lr·m̂/(√v̂ + ε/c)` — the `c` cancels except through `ε`. Residue: `c` varies per step so cancellation across accumulated moments is imperfect; plus the first steps before moments equilibrate. | +| **Muon** | exact | **none** for 2D params. Newton-Schulz maps `B = UΣVᵀ -> UVᵀ`, discarding the singular values, and the update magnitude comes from a shape-dependent `get_muon_scale_factor(size[0], size[1])`. Scaling `B` by `c` leaves `UVᵀ` unchanged. | + +Megatron's default is Adam, so for most users the weight-level effect is small — but +the reported norm is wrong regardless, and an SGD configuration sees the full effect. + +## Which models are exposed + +For a standard transformer every linear is column- or row-parallel, so `is_parallel=True` +is correct and the bug cannot trigger. + +It requires an architecture that *deliberately replicates* weights carrying real +gradient energy. Multi-head Latent Attention is exactly that — the low-rank +down-projections are replicated because the KV latent is shared by every head (each TP +rank needs the whole latent), and because the down-projection feeds a column-parallel +up-projection that requires the complete input on every rank. + +Megatron's own MLA does this explicitly +(`megatron/core/transformer/multi_latent_attention.py`): + +```python +if submodules.linear_q_down_proj in [TELinear]: + q_down_proj_kwargs['parallel_mode'] = 'duplicated' +... +if submodules.linear_kv_down_proj in [TELinear]: + kv_down_proj_kwargs['parallel_mode'] = 'duplicated' +``` + +**Scope:** any Megatron MLA model takes this path. DeepSeek-V3 has been measured +directly (see Evidence) using Megatron's own `pretrain_gpt.py`, confirming the +behaviour is not specific to any downstream model implementation. DeepSeek-V2 and V3.2 +share the same spec path and are expected to behave identically, though they have not +been run. + +## Evidence + +DeepSeek-V4-Flash, 4 layers, TP=8, EP=8, PP=1, DP=1, bf16, fixed seed (deterministic, +identical across reruns). Two linear backends, to show the behaviour is not specific +to one implementation. + +| backend | | it 1 | it 2 | it 3 | it 4 | it 5 | +|---|---|---|---|---|---|---| +| Lumen | current | 3.763 | 3.837 | 3.720 | 3.593 | 3.748 | +| Lumen | fixed | 2.472 | 2.520 | 2.445 | 2.359 | 2.462 | +| Lumen | ratio | 1.522 | 1.523 | 1.521 | 1.523 | 1.522 | +| TE | current | 4.098 | 4.177 | 4.079 | 3.921 | 4.112 | +| TE | fixed | 2.634 | 2.688 | 2.623 | 2.520 | 2.644 | +| TE | ratio | 1.556 | 1.554 | 1.555 | 1.556 | 1.555 | + +The constant ratio across iterations is the signature of a counting error rather than a +numerical one. The two backends differ slightly only because the runs have different +random initialisation. + +Loss is unchanged over the same window — differences in the 4th–5th decimal with no +systematic drift — confirming the fix alters only the clipping scale, not the +mathematics: + +``` +Lumen current: 12.60278 12.59553 12.58319 12.58988 12.58122 +Lumen fixed: 12.60291 12.59596 12.58199 12.58854 12.58262 +``` + +### DeepSeek-V3 on Megatron's own MLA + +Megatron `pretrain_gpt.py`, `--multi-latent-attention --transformer-impl transformer_engine`, +6 layers (3 dense + 3 MoE, following V3's `first_k_dense_replace=3`), hidden 7168, +128 heads, `q_lora_rank=1536`, `kv_lora_rank=512`, 256 experts, TP=8, EP=8, bf16. +No third-party model code; the only variable is the one-line change to +`Linear.reset_parameters`. + +| | it 1 | it 2 | it 3 | it 4 | +|---|---|---|---|---| +| stock TE | 12.228 | 12.048 | 11.886 | 11.815 | +| fixed TE | 9.011 | 8.848 | 8.739 | 8.656 | +| ratio | 1.3570 | 1.3617 | 1.3601 | 1.3649 | + +Loss is again unchanged (13.19588 vs 13.19578, 13.20511 vs 13.20663, ...). + +## History + +Introduced in `044903374` (2023-02-10, "QKV parameters unfused path fixes and +optimization" #66), which added weight marking to `transformer_engine/pytorch/module.py`. +Verified by bisect: the preceding commit `78b4e9339` (2023-02-07) has no weight-marking +call. + +The initial code drop (`996ea169c`, 2022-09-27) marked **only biases**, which in +column-parallel layers genuinely are sharded. `parallel_mode: Optional[str] = None` was +already reachable by `5612ba784` (2022-10-04), so the non-parallel case existed when the +unconditional marking landed. + +Everything since is refactoring, not semantic change: `c6a4a4e08` (2023-05-09) moved the +code into `module/linear.py`; v1.3 moved it from `__init__` into `reset_parameters`. The +behaviour is unchanged in ~3.5 years and is present both in NVIDIA `main` and in +ROCm/TransformerEngine `dev`. + +## Fix + +Mark only genuinely sharded weights: + +```python +is_parallel=self.parallel_mode is not None, +``` + +This restores TE's own documented default. `LayerNormLinear`, `GroupedLinear` and +`LayerNormMLP` should be checked for the same pattern. + +A downstream workaround needing no TE change is to clear the attribute after +construction, before the optimizer is built — the distributed optimizer copies TP +attributes onto its shards via `copy_tensor_model_parallel_attributes`, so the cleared +value propagates: + +```python +for param in replicated_module.parameters(): + param.tensor_model_parallel = False + param.partition_dim = -1 + param.partition_stride = 1 +``` + +## Caveats + +* Measured on DeepSeek-V4-Flash and DeepSeek-V3, both at TP=8, DP=1, with Adam. The + ratios are specific to those configurations; they follow `sqrt(1 + (TP-1)f)` and vary + with model, layer count and TP size. +* Both runs are depth-reduced (4 and 6 layers) for single-node capacity, so the + dense/MoE balance differs from production — real DeepSeek-V3 is 58 MoE + 3 dense, + which changes `f` and hence the ratio. Production TP is typically higher than 8, where + the scaling law predicts a larger factor. +* Verified that the reported norm and the clipping coefficient change. The + optimizer-sensitivity table is derived from the update rules, not measured end to end; + a convergence comparison against a reference curve has not been run. +* The DP>1 path is a code-path inference (attributes propagate via + `copy_tensor_model_parallel_attributes`), not measured. From f0c9f0e45add80dbc6309958097494838338fc2b Mon Sep 17 00:00:00 2001 From: wenchenvincent <32376000+wenchenvincent@users.noreply.github.com> Date: Tue, 25 Aug 2026 02:47:58 +0000 Subject: [PATCH 3/3] Apply the same fix to LayerNormLinear and GroupedLinear Both accept parallel_mode: Optional[str] = None and mark weights is_parallel=True unconditionally, exactly as Linear did (1 site in LayerNormLinear, 2 in GroupedLinear). LayerNormMLP is NOT affected: it has no parallel_mode parameter and is structurally column-then-row, so is_parallel=True is always correct there. The bias handling below each weight loop is already correct -- it branches on parallel_mode (row sets sequence_parallel, column marks the bias, None marks nothing). Only the weight marking ignored the mode. These two are latent rather than triggered: no Megatron wrapper instantiates them with parallel_mode=None today (TELayerNormColumnParallelLinear passes column; TEColumnParallelGroupedLinear/TERowParallelGroupedLinear pass column/row). They are reachable through TE public API, and fixing them keeps the three call sites consistent. Only the Linear fix is backed by measurement. --- docs/replicated_weight_gradnorm.md | 19 +++++++++++++++++-- .../pytorch/module/grouped_linear.py | 6 ++++-- .../pytorch/module/layernorm_linear.py | 3 ++- 3 files changed, 23 insertions(+), 5 deletions(-) diff --git a/docs/replicated_weight_gradnorm.md b/docs/replicated_weight_gradnorm.md index 5e09539860..c688922646 100644 --- a/docs/replicated_weight_gradnorm.md +++ b/docs/replicated_weight_gradnorm.md @@ -250,8 +250,23 @@ Mark only genuinely sharded weights: is_parallel=self.parallel_mode is not None, ``` -This restores TE's own documented default. `LayerNormLinear`, `GroupedLinear` and -`LayerNormMLP` should be checked for the same pattern. +This restores TE's own documented default. + +The same unconditional marking is present in `LayerNormLinear` (1 site) and +`GroupedLinear` (2 sites), both of which also accept `parallel_mode: Optional[str] = None`, +and is fixed the same way. `LayerNormMLP` is **not** affected: it has no `parallel_mode` +parameter and is structurally column-then-row, so `is_parallel=True` is always correct +there. + +Note the bias handling directly below each weight loop is already correct — it branches +on `parallel_mode` (row sets `sequence_parallel`, column marks the bias, `None` marks +nothing). Only the weight marking ignores the mode. + +No Megatron wrapper currently instantiates `LayerNormLinear` or `GroupedLinear` with +`parallel_mode=None` (`TELayerNormColumnParallelLinear` passes `"column"`; +`TEColumnParallelGroupedLinear` / `TERowParallelGroupedLinear` pass `"column"` / `"row"`), +so those two are latent — reachable through TE's public API but not triggered by Megatron +today. Only the `Linear` fix is backed by the measurements above. A downstream workaround needing no TE change is to clear the attribute after construction, before the optimizer is built — the distributed optimizer copies TP diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index f534da5c3b..ae636a398c 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -1617,7 +1617,8 @@ def set_tensor_parallel_attributes(self, defer_init=False) -> None: if grouped_weight is not None: set_tensor_model_parallel_attributes( tensor=grouped_weight, - is_parallel=True, + # Replicated when parallel_mode is None; see Linear.reset_parameters. + is_parallel=self.parallel_mode is not None, dim=1 if self.parallel_mode == "row" else 0, stride=1, ) @@ -1625,7 +1626,8 @@ def set_tensor_parallel_attributes(self, defer_init=False) -> None: for i in range(self.num_gemms): set_tensor_model_parallel_attributes( tensor=getattr(self, f"weight{i}"), - is_parallel=True, + # Replicated when parallel_mode is None. + is_parallel=self.parallel_mode is not None, dim=1 if self.parallel_mode == "row" else 0, stride=1, ) diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 479b346bfd..81c9784593 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -1703,7 +1703,8 @@ def reset_parameters(self, defer_init=False): for weight in self.weight_names: set_tensor_model_parallel_attributes( tensor=getattr(self, weight), - is_parallel=True, + # Replicated when parallel_mode is None; see Linear.reset_parameters. + is_parallel=self.parallel_mode is not None, dim=1 if self.parallel_mode == "row" else 0, stride=1, )