From 90bbf2ea7998479fe87305f9ab76fa3fe613bc20 Mon Sep 17 00:00:00 2001 From: wenchenvincent <32376000+wenchenvincent@users.noreply.github.com> Date: Mon, 31 Aug 2026 02:49:46 +0000 Subject: [PATCH 1/2] 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 56a9e191120fae6956b85ef1780c961278c24fb6 Mon Sep 17 00:00:00 2001 From: wenchenvincent <32376000+wenchenvincent@users.noreply.github.com> Date: Mon, 31 Aug 2026 02:49:47 +0000 Subject: [PATCH 2/2] 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. --- transformer_engine/pytorch/module/grouped_linear.py | 6 ++++-- transformer_engine/pytorch/module/layernorm_linear.py | 3 ++- 2 files changed, 6 insertions(+), 3 deletions(-) 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, )