Add an experimental DeepSeek-V4 TOP parity contract

[kaixih]

Summary

This draft adds a self-contained, experimental DeepSeek-V4 true-on-policy (TOP) numerical contract for the:

rollout -> SGLang prefill recompute -> Megatron training forward

path.

The main acceptance criterion is exact equality between the behavior-policy log-probabilities recomputed by SGLang prefill and the log-probabilities consumed by the Megatron training forward pass.

This is intentionally a draft/reference implementation. The change is large and is expected to be split into smaller, independently reviewable follow-up PRs before upstreaming.

Motivation

Matching the prefill scorer and trainer scorer requires more than synchronizing model weights or selecting nominally equivalent kernels. DeepSeek-V4's inference and training paths differed in floating-point association and implementation details across:

• RMS normalization and reduction trees
• attention reconstruction and TP reduction
• compressed attention and hyper-connections
• FP8 scale generation and grouped-GEMM accumulation
• routed and shared-expert MoE ordering
• full-vocabulary log-softmax precision and implementation

Small differences at these boundaries accumulate through the model and produce different per-token log-probabilities, even when both paths are individually deterministic.

What changed

Megatron/trainer path

• Added DeepSeek-V4 batch-invariant RMSNorm and fixed-tree reductions.
• Added deterministic TP reduction for the aligned attention path.
• Aligned query/KV normalization, value reconstruction, compressed attention, and hyper-connection arithmetic.
• Added trainer-compatible DeepSeek-V4 MoE routing, FP8 grouped-GEMM, routed scaling, combine ordering, and shared-expert accumulation.
• Scoped power-of-two FP8 scales to DeepSeek-V4 TOP.
• Added an FP32 true-on-policy full-vocabulary log-softmax option and batch-invariant log-softmax path, followed by the BF16 transport cast used by training.

SGLang/prefill path

• Added deterministic sparse-attention and MoE compatibility helpers.
• Added a version-pinned, fail-closed SGLang source arithmetic contract.
• Added runtime bootstrap propagation to spawned SGLang workers.
• Added explicit runtime/source contract switches with strict value validation.
• Added deterministic backend validation and fixed reduction/accumulation ordering for the tested TP8 path.

Contract and CI behavior

• Added argument validation for the new TOP log-probability contract.
• Kept the normal --ci-test TOP assertions enabled.
• Added a narrowly scoped option to disable only the unrelated post-sync weight checker, because DeepSeek-V4 has model-owned dynamic buffers that do not participate in that checker protocol.
• Rejected combining prefill recomputation with rollout routing replay because those options define conflicting scoring contracts.

End-to-end evidence

Validated with miles:latest on one 8xB200 node using the pruned four-layer DeepSeek-V4 model and TP8 for both SGLang and Megatron.

The test uses neutral sampling transforms:

temperature = 1
top_p = 1
top_k = -1

Positive run

• Five clean-cache prefill recomputations on the same SGLang engine.
• All 16 response-token FP32 log-probabilities were bitwise identical across all five passes.
• Prefill recompute versus Megatron training forward:
  • 128/128 active token/rank BF16 log-probabilities were exactly equal.

Contract-off negative run

The same branch and E2E were run with only the SGLang source numerical contract disabled. The workload still initialized successfully, completed rollout/prefill/training forward, and reached the intended parity assertion.

Result:

• 112/128 active token/rank entries mismatched.
• Mean absolute difference: 0.057136059.
• Maximum absolute difference: 0.25.

This distinguishes a real numerical-contract failure from a startup/configuration failure and demonstrates that the E2E does not pass without the alignment implementation.

Tests

• DeepSeek-V4 TP8 TOP E2E:
  • five-pass bitwise prefill repeatability
  • per-token/per-TP-rank BF16 prefill-versus-trainer exact equality
• Fail-closed source/runtime contract tests.
• Batch-invariant log-softmax gating tests.
• DeepSeek-V4 FP8 scale scoping tests.
• MoE request-local state restoration tests.
• Fixed-tree reduction invariance and autograd tests.
• Fused hyper-connection tests.
• Argument and CI checker behavior tests.

Final validation:

• 49 passed
• Black: 22 changed Python files clean
• git diff --check: clean

What this E2E proves

Within the tested contract, this is the key integration evidence for DeepSeek-V4 TOP: the SGLang prefill recompute scorer and Megatron training-forward scorer produce the exact BF16 per-token log-probabilities consumed by training.

It is analogous to the existing Qwen TOP E2E, with the additional requirement that five repeated DeepSeek-V4 prefill recomputations are themselves bitwise deterministic.

Current scope and limitations

This draft currently covers:

• the pruned four-layer DeepSeek-V4 checkpoint
• B200
• TP8
• the pinned SGLang version in miles:latest
• the tested 96-token E2E shape
• neutral sampling transforms
• forward/log-probability parity

It does not yet establish:

• arbitrary sequence lengths or batch sizes
• the full DeepSeek-V4 checkpoint
• other hardware or parallel configurations
• post-sampling behavior distributions for non-neutral temperature/top-k/top-p
• optimizer-update or full training-trajectory equivalence

Proposed follow-up decomposition

After reviewing this reference PR, the implementation can be split into smaller PRs such as:

1. Fixed-tree reduction and batch-invariant primitive tests.
2. TOP log-probability dtype and batch-invariant log-softmax support.
3. DeepSeek-V4 FP8 scale alignment.
4. Hyper-connection and compressor forward alignment.
5. DeepSeek-V4 MoE routing/GEMM/combine alignment.
6. Fail-closed SGLang runtime/source contract plumbing.
7. The final DeepSeek-V4 TP8 E2E.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.