Do not mark replicated weights as tensor-model-parallel

[wenchenvincent]

Summary

Linear.reset_parameters marks every weight is_parallel=True regardless of
parallel_mode. For parallel_mode=None the weight is replicated on every TP rank, so
Megatron's param_is_not_tensor_parallel_duplicate() — which gates get_grads_for_norm()
— admits it 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 and the reported gradient norm
is inflated.

Measured (MI355X, gfx950, TP=8, EP=8, bf16, fixed seed)

model	stock	fixed	ratio
DeepSeek-V4-Flash (4 layers)	3.763 / 3.837 / 3.720	2.472 / 2.520 / 2.445	1.522
DeepSeek-V3 (6 layers)	12.228 / 12.048 / 11.886	9.011 / 8.848 / 8.739	1.361

The ratio is constant per-iteration — the signature of a counting error, not a numerical
one. Loss is unchanged in both (differences in the 4th–5th decimal, no systematic drift),
confirming only the clipping scale moves, not the mathematics.

The DeepSeek-V3 run is Megatron's own pretrain_gpt.py with
--multi-latent-attention --transformer-impl transformer_engine and no third-party
model code — reproducible with upstream components alone.

Why this is architecture-specific

Inflation follows sqrt(1 + (TP-1)*f) where f is the replicated weights' share of the
true squared norm — gradient energy, not parameter count. Inverting the measurements
gives f ≈ 18.8% (V4-Flash) and 12.2% (V3), against a ~0.3% parameter share.

For a standard transformer every linear is column- or row-parallel, so is_parallel=True
is correct and the bug cannot trigger. It needs an architecture that deliberately
replicates weights carrying real gradient energy. MLA is exactly that — the KV latent is
shared by every head, and the down-projection feeds a column-parallel up-projection that
needs the complete input on every rank. Megatron's MLA passes parallel_mode='duplicated'
for q_down_proj/kv_down_proj, so DeepSeek-V2/V3/V3.2-family models take this path.

It also scales with TP: at f≈0.19 the law predicts 1.53x at TP=8 and 1.96x at TP=16.

Impact

Primarily diagnostic. The gradient norm is what practitioners use to judge training
health, tune --clip-grad, detect instability and compare against reference curves — and
it is wrong by 36–52% on these models.

Weight-level impact is optimizer-dependent, because Megatron's clipping is a global,
norm-based rescale that preserves direction and relative weighting:

optimizer	impact
SGD	full — gradients scaled down, and coupled weight decay (applied after clipping) becomes relatively stronger
Adam/AdamW	second-order — m -> c·m, v -> c²·v, so c cancels except through ε
Muon	none for 2D params — Newton-Schulz discards singular values

Megatron defaults to Adam, so most users see a wrong number rather than degraded
training. An SGD configuration sees the full effect.

Fix

is_parallel=self.parallel_mode is not None,

Restores TE's own default — _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS declares
tensor_model_parallel: False for a non-parallel tensor. TE never reads the attribute
back; it is set purely for downstream consumers.

LayerNormLinear, GroupedLinear and LayerNormMLP should be checked for the same
pattern — this PR changes only Linear.

History

Introduced upstream in NVIDIA/TransformerEngine 044903374 (2023-02-10, "QKV parameters
unfused path fixes and optimization" #66). Bisect: the preceding commit 78b4e9339 has no
weight-marking call. The initial code drop marked only biases, which in column-parallel
layers genuinely are sharded; parallel_mode: Optional[str] = None was already reachable
when the unconditional marking landed. Everything since is refactoring. Present in both
NVIDIA main and ROCm dev.

Caveats

• Both runs are depth-reduced (4 and 6 layers) for single-node capacity, so the dense/MoE
  balance differs from production; f and the ratio would shift at full depth.
• Both at TP=8, DP=1, with Adam. The DP>1 path is a code-path inference (attributes
  propagate via copy_tensor_model_parallel_attributes), not measured.
• The optimizer-sensitivity table is derived from the update rules, not measured end to
  end; no convergence comparison against a reference curve has been run.