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. Done: trainer-side fixed-tree reduction and batch-invariant DSV4 query RMS (#1912).
2. No separate Miles runtime PR needed: Megatron enable_batch_invariant_mode() globally registers the accelerator implementation of aten::_log_softmax, so the existing Miles torch.log_softmax scorer automatically uses the batch-invariant kernel. The test-only split #1964 was closed.
3. Partially implemented: DeepSeek-V4 FP8 scale alignment. Miles #1182 already selects SGLang per_block_cast_to_fp8 for [128, 128] trainer-to-rollout hot updates, producing FP32 power-of-two expert-weight scales. GB200 validation confirmed bitwise-equal logical scales and FP8 weight bytes against Trainer TE with force_pow_2_scales=True for representative weights. SGLang #33005 supplies the remaining Triton-MoE activation-scale propagation.
4. DSV4 attention preparation and reconstruction alignment: add the SGLang rollout/prefill fixed-tree query RMS path—the inference-side counterpart of [DSV4 TOP 1/N] Make trainer query RMS batch invariant #1912—and align Q/KV normalization, RoPE generation and application, KV-cache writes, inverse RoPE, the single-group value-reconstruction GEMM, and deterministic wo_b TP reduction.
5. DSV4 sparse-attention and indexer alignment: use a supported common forward kernel for trainer and prefill, and replace the hash_topk FP32 runtime monkeypatch with a supported integration.
6. DSV4 C4 compressor forward alignment: make the trainer use the SGLang-compatible forward with correct autograd support.
7. DSV4 hyper-connection alignment: use common HC-head, HC-pre, and HC-post forward kernels with trainer backward support.
8. Routed and shared MoE expert-compute alignment: align FC1, activation, FC2, and FP8 GEMM accumulation through a supported SGLang-compatible trainer path, without rerunning TE grouped GEMM over Triton results.
9. MoE routing, combine, and collective alignment: align routing-weight cast and application timing, stable expert/token accumulation order, shared-expert TP reduction, and routed/shared addition order.
10. Remove the diagnostic source-patcher plumbing: delete the YAML/bootstrap path after every arithmetic boundary above has a supported integration, while retaining fail-closed runtime contract validation.
11. Final DeepSeek-V4 TP8 E2E: run the parity test without DUMPER_SOURCE_PATCHER_CONFIG.

[gemini-code-assist]

Caution

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

[kaixih]

Correction after GB200 validation: the Miles trainer-to-rollout expert-weight path was already aligned by Miles #1182. For [128, 128] block FP8, it uses SGLang per_block_cast_to_fp8, which returns FP32 power-of-two scales.

RL hot-update path

State	Weight scales	Activation scales
Before SGLang #33005	Power-of-two	Continuous
After #33005	Power-of-two	Power-of-two

Pure SGLang inference with SGLANG_DSV4_FP4_DEQUANT=1

State	Weight scales	Activation scales
Before #33005	Power-of-two	Continuous
After #33005	Power-of-two	Power-of-two

On GB200 with radixark/miles:latest, the public Triton-expert hot-update quantizer produced power-of-two scales and matched Trainer TE (force_pow_2_scales=True) bitwise in both logical scales and FP8 weight bytes for representative weights and scale-boundary cases. No additional Miles weight-scale feature PR is needed; #33005 is the remaining activation-side change.