Mamba linear CE parity deep dive

An output path can preserve tensor shapes while still drifting on the logits-to-loss contract. That is why a CE parity reproducer is worth publishing. It narrows the question to the only thing that matters: does the alternate path preserve the same logits-to-loss behavior?

The near-copy version keeps the output-layer contract visible: one path uses a fused linear-plus-cross-entropy module and the other keeps a plain column-parallel output layer until the loss-path contract is restored. That is closer to the real failure than a generic cross_entropy(hidden @ W.T, targets) toy. This makes the post the output-side sibling of Porting to Megatron friction: the useful work is in keeping the seam visible rather than pretending shape parity proves enough.

What the checked-in examples prove

If you want the local proof surfaces before the article prose, start with MegaCpp example index, Mamba linear CE parity sample, Mamba linear CE parity near-copy, and Megatron FLCE Hopper near-copy. The compact sample keeps the logits-to-CE interface tiny enough to inspect directly. The checked-in sample is more useful when you care about the real integration seam because it keeps the class-level boundary visible: one path owns a fused output-and-loss surface, the other exposes a plain column-parallel layer until the patch restores the expected contract.

Why shape parity is not enough

Shape checks are necessary and still too weak. Two paths can agree on output dimensions and still drift if one accumulates its running max and denominator in a different dtype, or if one loss path averages and reduces tokens at a different point than the other. That is why CE parity here is a behavioral property, not a shape property. The real question is whether the runtime still consumes the same logits-to-loss boundary it was designed for.

The distributed case is stricter still. Once the output layer is vocab-parallel or sequence-parallel, parity is no longer only a local logits question. A serious parity proof has to name where the max reduction happens, where the denominator sum reduction happens, and whether the final loss is normalized by local tokens or by the global token count. On sequence-parallel lanes it also helps to state the layout explicitly, usually [S, B, V/TP] or [S/SP, B, V/TP], because "parity passed" otherwise often means only that the local tensor shapes looked plausible.

That accumulator story deserves to be named plainly. In vocab-parallel or chunked CE, the running max and running denominator are part of the contract, not an invisible implementation detail. If one path upcasts those accumulators and another keeps them in the execution dtype, the loss can drift before the final token reduction ever happens. That is why a useful parity note says where accumulator ownership lives: inside a fused kernel, across TPQuick term guideTPTensor parallelism splits each linear's weights (QKV, O, MLP gate/up/down) across GPUs. On 8× H200 with TP=8 each GPU owns 1/8 of every matmul's columns or rows, so one big matmul becomes 8 smaller ones that all-reduce at the layer boundary. Cost: one all-reduce per attention and per MLP — heavy bandwidth, so TP is usually bound to a single NVLink/NVSwitch island (1 node of up to 8 GPUs). Embeddings, layernorms, and optimizer state stay replicated across the TP GPUs. Use TP when a single layer's weights don't fit on one GPU, not to scale past one node.GroundingAbout: parallelism map overview Example: TP partition-shape sample Reference: tensor parallel and sharding reductions, or only after a materialized-logits fallback. Megatron FLCE on Hopper is the closest sibling when the question becomes how that boundary is implemented on the fused side.

Sequence-parallel lanes add one more place to be explicit. Once the tensor is [S/SP, B, V/TP], the parity claim should also say whether the final mean divides by local visible tokens or by the global non-padding token count after cross-rank reductions. Without that statement, a reproducer can pass the obvious shape checks and still change gradient scale across ranks, which is exactly the kind of seam drift that later gets misread as a generic training-health problem rather than an output-layer problem. Liger FLCE reduction none is the neighboring post when the reduction contract itself is the thing under test.

The easiest concrete example is small enough to keep in your head. If one lane averages over 1024 visible local tokens while a TPQuick term guideTPTensor parallelism splits each linear's weights (QKV, O, MLP gate/up/down) across GPUs. On 8× H200 with TP=8 each GPU owns 1/8 of every matmul's columns or rows, so one big matmul becomes 8 smaller ones that all-reduce at the layer boundary. Cost: one all-reduce per attention and per MLP — heavy bandwidth, so TP is usually bound to a single NVLink/NVSwitch island (1 node of up to 8 GPUs). Embeddings, layernorms, and optimizer state stay replicated across the TP GPUs. Use TP when a single layer's weights don't fit on one GPU, not to scale past one node.GroundingAbout: parallelism map overview Example: TP partition-shape sample Reference: tensor parallel and sharding=8 fused lane divides only after cross-rank reductions over 8192 non-padding tokens, the per-token logits can still look fine while the returned gradient scale is different by 8x. That is why this post keeps reduction ownership next to Megatron FLCE on Hopper instead of treating it as an implementation footnote.

That is also why the near-copy reproducer matters more than a toy script. A toy hidden @ W.T check can validate arithmetic in isolation, but it cannot show whether a path swap quietly changed the owned output-and-loss seam. The checked-in near-copy files keep that seam visible enough to debug, which is the same operational reason this article belongs next to Megatron FLCE on Hopper, Liger FLCE reduction none, Author Mamba3 spec inside Megatron, and Loss curves and divergence playbook.