Silent numerical-correctness bugs in Megatron Core parallelism composition

Human summary: I did a fuzz of Megatron's codebase with Astra and spmd_types looking for distributed parallelism bugs, and it turned up these. I haven't personally verified every single one, but some of these overlap with existing bug reports, and I also ran this experiment on our internal training codebase and every single bug Astra reported was real, so.... I'm expecting most of these are real (even if they are potentially in configuration permutations you might not care about. Some of these bugs are definitely live for Megatron-Bridge configs though!)

There's also a few non-spmd_types numerical correctness bugs that are reported at the very bottom that Astra also noticed, but they're probably real bugs too.

I did some work on preparing this report and skimmed it, but the rest of this is all AI generated. It usually took me a few hours to prepare and human review fixes for each individual bug in the internal codebase.

AI report

These bug reports are off of:

• Megatron-LM origin/main at 7a9e59de731bd4a470bf89c577004fbc51345a88 (2026-09-17)
• Megatron-Bridge main at 34d8d4b8281f4ef3955f42c3d9d11eaf445a2434 (2026-09-17), used only for configuration/reachability evidence and for the one Bridge-side bug in Appendix B

Each finding has a self-contained CPU reproduction (two or four Gloo ranks, posted as collapsed blocks in the comments below) that calls the actual Megatron functions involved and replaces only CUDA-only kernels or allocations with torch equivalents. The reproduction prints the observed value, the reference value, and where applicable a corrected control.

Each reproduction also states the violated dataflow contract in the notation of spmd_types, which is how these bugs were found. Per process group, a tensor is one of:

• R replicated: every rank holds the same full value.
• P partial: every rank holds a partial sum; the true value is the SUM over the group.
• S(d) sharded along dimension d; I/V invariant/varying scalars.

The rules the bugs break are few: a P value may not enter a nonlinear op (such as a division) before it is reduced; treating a P value as R (keeping one rank's copy) or an R value as P (summing the copies) is a type error; a collective must be issued on the group whose axis the type refers to; and consumers that require S(d) must not receive R. Each finding's checker: line is the one-line rejection spmd_types emits for the buggy operation.

None of these raise an error. Shapes stay valid and training continues with wrong losses, wrong gradients, or wrong activations.

Summary

Appendix A lists five findings from the same audits that are already fixed on main. Appendix B has one Megatron-Bridge bug. Appendix C lists candidates that were confirmed numerically but not counted because they are constant scale errors without a dataflow typing witness.

1. Vocab-parallel label smoothing uses the local vocabulary

Location (7a9e59de7): megatron/core/tensor_parallel/cross_entropy.py:158 sets vocab_size = exp_logits.size(-1), which is the per-rank partition size. Line 175 computes log_probs.mean(dim=-1) over the local shard only. Line 178 saves the local size into ctx.vocab_size, and backward reuses it at line 191.

Trigger: TP>1 and label_smoothing>0 through vocab_parallel_cross_entropy. The base CE correctly all-reduces the max, target logit and exp-sum; only the smoothing tail is local.

Effect (TP=2, smoothing 0.2, two positions): rank losses [1.4401896, 1.0951819] versus [0.6401896, 0.9951819]; the global-vocabulary formula gives [0.8401898, 0.8618487]. Maximum gradient error 0.1333.

Contract violated: the local mean over the sharded vocabulary dimension is a partial value; it is combined with the already complete CE loss without a reduction.

Fix: sum the local log-probability sums over the TP group and divide by the global vocabulary size; use the global size in the smoothing coefficient and in backward. PR #5522 proposed exactly this and was closed without merging. A comment posted on issue #737 on 2026-09-17 by another user independently reports the same analysis.

2. Gathered logits are passed to vocab-parallel cross entropy

Location: megatron/core/models/gpt/gpt_model.py:292 builds the output layer with gather_output=not self.parallel_output; runtime_gather_output (line 579/671) can also force a gather. When labels are supplied, the logits go to compute_language_model_loss, which always calls a vocab-parallel CE (megatron/core/models/common/language_module/language_module.py:127).

Trigger: TP>1, labels supplied, and either parallel_output=False or runtime_gather_output=True. Both are public model options.

Effect (TP=2, logits [0,1,2,3], label 3): observed loss 1.1333369 versus correct 0.4401897, i.e. correct plus log(2). Physical gradient [0.0160, 0.0436, 0.1184, 0.3220] versus reference [0.0321, 0.0871, 0.2369, -0.3561]; the target term is lost.

Contract violated: the CE treats the last dimension as one vocabulary shard. Gathered logits are replicated over TP, not sharded on the last dimension. The TE CE path already asserts this contract; the unfused path does not.

Fix: assert sharded logits when computing parallel CE, fall back to ordinary CE when the logits were gathered, or reject the combination.

3. MTP per-token loss uses local token-count ratios

Location: megatron/core/transformer/multi_token_prediction.py:1113 (original_num_tokens = loss_mask.sum()) and line 1216, which scales each MTP loss by original_num_tokens / num_tokens_safe. Both counts are local to the DP/CP rank. The later global normalization in finalize_model_grads.py divides the summed gradient by the global main-token count and cannot undo a nonlinear local ratio.

Trigger: MTP with calculate_per_token_loss=True and uneven valid-token counts across DP or CP ranks (any SFT-style masking).

Effect (two DP workers with main/MTP counts 4/3 and 2/1, unscaled MTP gradient sums 6 and 10): implemented gradient 4.6666665, single-worker reference 4.0.

Fix: all-reduce both token counts over the same DP/CP group used for gradient normalization before forming the ratio.

Related: #4896 reports the same ratio-averaging pattern in MTP loss logging and states gradients are unaffected; this report shows the gradient path has the same problem. #3943 and #1532 were earlier MTP scaling bugs, now closed.

4. MTP full recomputation uses the global TP group

Location: multi_token_prediction.py:1872 passes parallel_state.get_tensor_model_parallel_group() into the TE checkpoint path even though the layer stores its explicit group as self.tp_group. The non-TE branch has the same problem through tensor_parallel/random.py:624 and :649, which call the split/gather helpers with no group and fall back to MPU state. TransformerBlock passes self.pg_collection.tp in the corresponding path (recompute.py).

Trigger: MTP with full recompute and distributed saved activations, and an explicit ProcessGroupCollection whose TP group differs from the global MPU TP group.

Effect (four ranks; explicit TP groups [0,2],[1,3], MPU groups [0,1],[2,3]): rank 0 restores [0,1,12,13] instead of [0,1,2,3]. Same shape, mixed data.

Fix: pass self.tp_group to the TE path; let tensor_parallel.checkpoint() and the 1-D split/gather helpers accept an explicit group. #7403 (merged 2026-09-16) fixed the analogous omission for the MTP input projections but not this path.

5. GPT scatters the MoE padding mask over the wrong TP group

Location: gpt_model.py:369. The decoder input is scattered with group=self.pg_collection.tp just above; the padding mask scatter omits the group and uses the global MPU TP group.

Trigger: MoE with padding_mask, SP, explicit pg_collection.tp different from global TP.

Effect (four ranks, seq len 4): ranks 1 and 2 receive mask [False, True] where [True, False] matches their hidden states. Router z-loss on rank 1 changes from 121.0004 to 100.0009; rank 2 from 9.2939 to 4.5238. One real token is excluded and one padding token included.

Fix: group=self.pg_collection.tp on the mask scatter. Open PR #5088 and closed PR #5563 refactor the same code (PP-stage coverage) and #5563 also passes the explicit group.

6. Cross-attention Q/KV projections omit the explicit TP group

Location: megatron/core/transformer/attention.py, CrossAttention.__init__, builds linear_q and linear_kv without tp_group or pg_collection. Base Attention and SelfAttention pass both to their projections. Reachable from the stock T5 decoder specs, which accept a ProcessGroupCollection.

Trigger: cross attention with SP and an explicit TP group different from global TP.

Effect (four ranks): the SP all-gather reconstructs [0,1,10,11] on ranks 0/1 and [2,3,12,13] on ranks 2/3 instead of [0,1,2,3] (sample A) and [10,11,12,13] (sample B). Each sequence contains two different data replicas.

Fix: pass tp_group=self.pg_collection.tp and pg_collection=self.pg_collection to both constructors.

7. LLaVA scatters combined embeddings over the wrong TP group

Location: megatron/core/models/multimodal/llava_model.py:1103. LLaVAModel builds the inner language model with scatter_embedding_sequence_parallel=False and owns the one required scatter, but performs it without a group.

Trigger: LLaVA, SP, explicit pg_collection.tp different from global TP.

Effect (four ranks): after the following gather on the explicit group, sample A [0,1,2,3] becomes [0,1,0,1], sample B [10,11,12,13] becomes [12,13,12,13].

Fix: group=self.pg_collection.tp on the scatter.

8. Learnable attention sinks are excluded from the clipping norm

Location: megatron/core/transformer/dot_product_attention.py:129 creates softmax_offset with one entry per local head, so it is head-sharded across TP, but never marks it tensor_model_parallel. TEDotProductAttention with TE 2.15 leaves the flag unset too. megatron/core/optimizer/optimizer.py:299 then applies param_is_not_tensor_parallel_duplicate, which keeps the parameter only on TP rank 0.

Trigger: softmax_type="learnable" (GPT-OSS) with TP>1 and gradient clipping. The Bridge GPT-OSS GB200 performance recipe uses TP=2, CP=4.

Effect: sink gradients 3 and 4 on two ranks give norm 3 instead of 5; with clip_grad=4, clipping is skipped (scale 1.0 instead of 0.8). In a full small attention block with real gradients: norm 1.22608447 instead of 1.22774088, next-output error 2.4e-4 after one clipped step. Affects the norm used to clip all parameters.

Fix: set tensor_model_parallel=True (and partition_dim=0) on softmax_offset. Same class of error as #5916 (expert grads undercounted, fixed 2026-07-25) on a different parameter.

9. CP-local loss means are averaged instead of the global mean

Location: megatron/core/pipeline_parallel/schedules.py:295 (forward_step_calc_loss) divides each CP rank's masked loss sum by its own valid count when calculate_per_token_loss=False. pretrain_gpt.py returns CP-local sums and counts. DDP then averages gradients across DP×CP.

Trigger: ordinary GPT training with --context-parallel-size 2 --eod-mask-loss (per-token loss defaults to false), or any loss masking with unequal counts per CP piece. Wan's real-data path pads to 2*CP and masks padding, with CP=4/8 recipes.

Effect (8 tokens, CP counts 4 and 1, 4-weight linear predictor):

observed:  [ 0.83685845, -0.08999243,  0.29482320, -0.10633019 ]
reference: [ 0.51100414, -0.14398789,  0.05773243, -0.17012830 ]

It computes (L0/n0 + L1/n1)/2 rather than (L0+L1)/(n0+n1), so the gradient direction changes, not just its scale.

Fix: sum the valid counts over CP before dividing (as the per-token path does).

10. Router local token count is overwritten by the global count

Location: megatron/core/transformer/moe/moe_utils.py:269 (get_tokens_per_expert_and_token_count) passes local_tokens_per_expert to reduce_from_tensor_model_parallel_region. _reduce in tensor_parallel/mappings.py mutates a contiguous input in place (the comment says so). With a padding mask, line 283 then derives local_num_tokens from the mutated tensor, so both returned counts are global. router.py:452/509/555 passes this "local" count as valid_token_count, and the aux loss is multiplied by it and by the TP/CP group size.

Trigger: MoE aux loss with padding_mask (sequence packing from the data scheduler), per-token loss, TP or CP > 1. Any non-null mask selects the branch.

Effect (TP=2, 5 valid of 8 tokens): both ranks get local_num_tokens=5. Router gradient norm 2.2258655 instead of 1.1129328, i.e. exactly TP× too large.

Fix: clone before the in-place reduce, or compute local_num_tokens before calling the collective. The 2026-09 change from / to // (#7081) did not affect this.

11. GDN shared output-normalization weight misses the TP gradient sum

Location: megatron/core/ssm/gated_delta_net/common.py:267 builds out_norm over value_head_dim, one vector applied to all heads, while heads are split across TP. megatron/core/distributed/finalize_model_grads.py:452 sums such gradients only when SP is enabled or the parameter name contains q_layernorm/k_layernorm.

Trigger: GDN with TP>1 and SP disabled. The Bridge Qwen3.5 full-SFT recipes use TP=2/4 without SP and MCore's GDN.

Effect (small full GDN module, TP=2): all other parameter gradients match the unsharded reference to 7.5e-9; out_norm gradients are [0.1200, 0.0400] and [0.0197, -0.0077] on the two copies versus the complete [0.1397, 0.0323]. After 20 Adam steps the two copies differ by 0.0039; adding the TP sum keeps every step within 2.4e-7.

Fix: mark out_norm.weight for TP gradient summation regardless of SP (or set sequence_parallel on it as TENorm does for SP).

12. Megatron-FSDP checkpoint restore u