Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Train-inference consistency workflow: C0–C10
Turn a train-inference difference into a repeatable experiment with fixed input and fixed weights, locate the first divergence, fix or optimize, then return to the real training loop. Keep the user's model, precision, parallelism and consistency target. Switches, thresholds, passing scope and implementation paths from historical cases, reference PRs and implementations already in the training framework (such as an existing true-on-policy path) are experience, not defaults for the current model, dataset or hardware; measure their numerics and training-step time before reusing them.
Forward consistency, backward conformance to the gradient contract, training quality and performance are accepted separately; none replaces another. User config comes first. Recover missing information from the source and resolved config, and ask only for what cannot be recovered and affects the approach. Choose the low-precision recipe by hardware and existing paths, not by a default precision, model or framework. Changing the model or precision requires rechecking coverage and acceptance evidence.
Work in one directory per task, with state changes in an event log and artifacts carrying provenance sidecars. Details are in validation.md (entry gate table, dual-path contract, tiered criteria, SFT three-way control, two-step RL sequence) and numerical contracts.
Archive materials as you go: each run ID links config/scripts, execution identity, raw results and conclusion; the step index states the current candidate, failed/excluded items, unverified items and the human-confirmed version. C1 keeps the full tuning comparison table; C10 references the reviewed baseline ID. Layout: archive experiment materials by step.
Gate table
Iterate as "single operator / fixed replay -> proxy -> staged full model"; do not rerun every gate each time, and do not skip the human stop points at C1 and C9. Inference-only tasks apply only the relevant gates and state that the training loop is not accepted.
| Gate | Input | Action | Pass condition | Artifacts | See |
|---|---|---|---|---|---|
| C0 Module inventory and route | checkpoint, vLLM implementation, target (BI=1 / BI=0), hardware | Per module: inference kernel, whether training implements the same math, quantization rounding, discrete choices; choose FSDP2 / Megatron-Bridge / mLite | Inventory covers all modules; route fits "BI=1 needed? how many backwards missing? what parallelism?" | Module inventory, route and rationale | Training-side model support |
| C1 Recipe and strong baseline | Model, target context length, hardware | Decide weight/activation/KV precision, quantization rules, runner, parallelism, execution config; decide the training master start (BF16 original or dequantized weights; do θ0 bytes equal the official checkpoint); agree scoring scope, consistency standard, performance budget (inference throughput, and limits for training forward and forward+backward against the native training path on the main topologies; a human decides which parallelism to enable and which topologies are main); tune BI=0 with vLLM Rust Bench at 8192 in / 1024 out / ignore EOS / c=1–64; for MoE, measure MoE / all-to-all backend combinations | Baseline = the formal recipe's BI=0 config (including R3, PC, gpu_mem), not the serving-only optimum; multi-round measurements and backends recorded; freeze the candidate and submit for human audit | Tuning table, resolved config, image digest, audit materials | Precision, quantization and QAT, Recipe and evaluation, Performance |
| C2 Proxy model | Full model structure | About 4 layers covering all operator types, needed adjacent-layer interactions and state paths; real hidden/head/expert/quantization geometry and needed parallelism constraints; coverage table "full-model feature -> proxy layer/module -> test case" | Uncovered items listed; the proxy is the daily development environment | Coverage table, proxy config | Training-side model support |
| C3 Inference BI/PD | Proxy, fixed weights and tokens | Verify batch and prefill/decode invariance separately, covering chunks, mixed scheduling, cache state; consider kernel suitability for training / large token counts. On failure, coarse-scan layer boundaries, bisect, capture identical operator inputs, fix kernel/metadata/layout/state | Bitwise identical in all three cases under fixed input; proves only consistency inside inference | Comparison matrix, minimal failing replay | BI kernel appendix table, Fault localization; before changing inference read implementation ownership; sync source per runtime provenance (same-commit worktree + per-file hash) |
| C4 Structural performance | Proxy, real operator shapes | Profile prefill/decode (and training forward and forward+backward once wired): launches vs tokens, CPU gaps, sync, repeated quantization/copies/allocation, abnormal workspace; target CUDA Graph captures/replays correctly | No per-token Python loop or many-small-kernels defects (autoregressive decode steps do not count); consistency retest passes after the fix | Profile, defect list | Performance |
| C5 Training forward↔inference | Proxy, fixed weights, input/output tokens, state | Reuse inference primitives and quantization, layout, mask, position, routing, reduction, logprob semantics; keep the training parameter container and native entry point; same-packing A/A and different-packing A/C first; then compare native old-logprob recomputation and the grad-enabled actor forward separately per the dual-path contract | Both paths: valid logprob mismatch=0 in original dtype; record non-empty valid token count and execution identity. Inference BI/PD, single operator or diagnostic substitute paths do not replace this gate | Dual-path receipt | Training-side model support |
| C6 backward / STE | Forward implementation from C5 | Reuse existing backward, autograd, communication, recompute; where the gradient passes quantization rounding, decide the STE surrogate (identity/clipping, master or dequantized value, scale rule). A kernel without autograd does not force an approximate gradient; no AST rewriting or monkeypatching in the production path | Every module has a "reuse / glue / newly written" verdict and rationale, and its backward uses only what the forward saved; the training step is within budget on the main topologies fixed at C1 | Module-level backward inventory, training-step time table for the main topologies | Backward validation, STE contract |
| C7 Independent gradient validation / SFT control | C6 implementation, same input, weights, saved state, same-dtype upstream gradient | Compare against a same-contract reference that does not call the backward under test: dX, dW, other trainable gradients, one optimizer step; fix tolerances first; report bitwise and tolerance separately. When required, compare production, an independent surrogate with the same visible forward, and native training per the three-way SFT control | Independent comparison passes; without SFT, mark "quality not covered", a loss decrease is not quality. Forward 0 diff or finite gradients do not replace this gate | Gradient audit, three-way SFT results | Same as above |
| C8 Proxy RL loop | Proxy, real framework entry point | Follow the exact sequence: θ0 compare -> update to θ1 -> θ1 compare -> update to θ2 -> θ2 compare; check finite, meaningful gradients, actual master change, normal exit | No cross-version comparison; a forward skipped by zero advantage does not count; proxy limited to logprob alignment: mark "loop not verified"; proxy asked to verify forward only: full-model zero-update comparison first, then the training gates; never start multi-step RL directly | Two-step RL receipt | Two-step RL sequence |
| C9 Recipe review and full-model entry gate | Launch script, image digest, code version, full resolved config, differences from the reference recipe (verl #7688 etc.), proxy smoke | Human review; start the full model only after the recipe version is confirmed (data, length, batch, parallelism, precision/quantization, optimizer); fill in the entry gate table; then add layers, length, concurrency, parallelism one at a time and verify mixed scheduling, cache/Graph, multi-step loop; with checkpoint restore in scope, check alignment after the first restore | No escalation while C5 is fail / inconclusive / without artifact; two-step acceptance is the formal recipe with only fewer steps, same sequence/batch/data/algorithm; on full-model failure, shrink to a minimal replay first | Review materials, entry gate record, two-step results | Checkpoint, resume and memory, recipe review materials and two-step agreement |
| C10 Optimization and final E2E | Full-model consistency passed, delivery PRs, reviewed recipe | E2E profile -> BI=1 hotspots (shape, dtype, layout, frequency, time share) -> representative micro bench -> correctness and performance on all target shapes -> E2E. Triton is forbidden; use CuTe DSL or better, CUDA/CUTLASS, or an existing implementation. Each round regresses inference, training forward and gradients by impact | On the delivery image through the real trainer entry: θ0 compare, update, θ1 compare, update, θ2 compare; 0 diff = equal original-dtype logprob for the same weight version, token history and valid mask; 8K/1K Rust Bench BI=0/BI=1 per the pairing protocol; no added long runs or gates | Delivery manifest, pairing results, loop receipt | Performance, Paired acceptance, Delivery |
The consistency implementation uses mLite by default: it separates model, operators and runtime, so the inference forward can be reused, backward/STE wired in independently, and the existing verl_mlite interface and MCore distributed optimizer wrapper kept. The training baseline follows the default MCore path of the existing verl recipe without the alignment implementation, and is accounted for separately from C1's vLLM BI=0 inference baseline. mLite does not inherently guarantee 0 diff or better performance; verify the target model's parallelism, recompute and checkpoint support with the actual recipe. Both arms share one recipe and one image, and VLLM_BATCH_INVARIANT selects the arm (see bi-switch-recipe); the delivery contract is in validation.md.
Run the cheapest failure criterion first
- host metadata / shape arithmetic: minimal counterexample of a pure function and boundary parameterization.
- Single operator: same input, real dtype/layout, non-finite detection, repeat / cross-process / Graph.
- frozen batch replay: the original failing A/C and A/A, full valid-token comparison.
- engine: batch composition, prefill/decode, chunk boundaries, cache/Graph, actual parallelism scope.
- Real training: non-degenerate reward/advantage, finite non-zero gradients, optimizer step, parameter change, frozen-token alignment after the update.
- Production acceptance: full model, agreed multiple steps, restore lifecycle, quality evaluation and paired performance.
If a second run adds no localization information, change the probe. Not reproduced, empty artifacts and an unexecuted path are not passes. When the failure is in backward, replaying the forward is not enough: save checkpoint, batch, RNG and state and move to a minimal backward experiment. Details are in validation.md.
RL delivery gates
Training and inference use the same immutable image digest and the same Python/Torch and runtime dependencies. They may run on separate nodes and processes, but two environments inside one image cannot be used to dodge compatibility. Complete at least two real updates through the target RL framework's actual entry point and verify each round: original-dtype logprob 0 diff before the update, an actual parameter update, and 0 diff for the new policy after it. Saved rollouts, offline replay, separate images and cold start are diagnostics only. This is a delivery agreement, not a technical claim about all RL systems.
Deliver one vLLM PR, one Megatron PR, one verl PR and one unified image, targeting vllm-project/vllm:main, NVIDIA/Megatron-LM:dev and verl-project/verl:main. Non-observability changes in verl should be limited to the recipe and the image script where possible. Follow from clean PR to delivery: check the target branch and full diff, pin cross-repository commits, build the image from clean source, and fill in the delivery manifest.
Baseline and consistency run as two parallel tracks
- Baseline track: C1's BI=0 tuning, backend combinations and multi-round performance evidence. Do not sacrifice default performance for consistency in advance; after reasonable tuning, freeze the candidate and wait for human audit.
- Consistency track: do not wait for the baseline; advance through C2–C8 on the proxy. C9 still requires the entry gate; parallel progress is not a reason to go to the full model early.
- Share the model/precision/hardware/image contract; isolate code, environment and artifacts; keep consistency patches out of BI=0. While the baseline is not final, mark the scope consistency covers and re-verify after config changes.
- Do not add nodes beyond the budget. On a single node, GPU experiments time-share; during a formal benchmark, run no GPU, heavy CPU or I/O work from the other track. The final performance gap is accepted against the human-confirmed strong baseline with full consistency evidence.
Two human stop points
After C1: baseline audit. Once measurements support "reasonably tuned", freeze the candidate, report and wait; do not keep chasing the performance limit. Materials: model/weight version, hardware and node count, image digest / build time, code and dependency versions, reproducible launch and Rust Bench commands; weight/activation/KV precision, runner, parallelism and scheduling, Graph/cache state, actual operator backends, and a statement that alignment patches do not change BI=0; configs tried and why, valid tokens per concurrency, warmup, repeat count, throughput/latency and variation, failed items and uncovered scope; the config proposed for locking, the statistical measurement basis, repeat count, per-concurrency criteria and budget, and open questions for humans. An unconfirmed candidate baseline must not be the final denominator. Re-audit after changing model/precision/hardware/image; while waiting, hand over or release nodes as needed.
Before C9: recipe review. See the gate table. Changes to workload or numerical/update semantics are reviewed again; pure path or log fixes are not.
Find by problem
| Current problem | Read first |
|---|---|
| Setting the consistency target or evidence basis | Experiment contract |
| Reproducible input not yet saved | artifacts, rollout-engine |
| rollout and old logprob differ: side unknown / stable failing sample / identical operator inputs still differ | Fault localization, numerical contracts |
| Results change with packing, CP/TP/EP | Layout and length |
| Fails only under Graph, compile or cache | Execution state |
| Wiring in backward/STE, comparing gradients, SFT | Reuse backward and backward validation (three layers) |
| Going to the full model, shrinking a full-model failure, declaring completion | validation.md |
| Candidate/backend/image changed: is old evidence still valid | Candidate evidence ledger, inferences no longer carried over |
| Aligned; recovering performance, profiling, ablation, BI kernels | Performance, BI kernel appendix table |
| Separating quality, fixed-length Rust Bench and normal exit | Evaluation and paired acceptance |
| Image, source, wheel, submodule, overlay versions confused | Runtime provenance |
| Three-repo ownership, PRs, unified image | Three-repository delivery |
| GB200 devices, NVSHMEM, DeepEP, cross-node initialization failure | Communication fault ladder |
checkpoint save/load OOM, unexpected pos, host memory rising, first resume misaligned |
Checkpoint, resume and memory |
| BI=0 baseline arm, two modes in the same recipe | bi-switch-recipe |
| Requantization rules, QAT selection, activation fake quantization | Precision, quantization and QAT |
| Many repeated samples within a rollout group | rollout-engine |
Different seeds give the same samples (greedy n repeats are identical; SGLang deterministic mode without a seed always uses sampling_seed=42) |
rollout-engine §1 |
The ten kinds of evidence that must not be confused are in evidence and experiment rules §2. Writing comparisons and replays: save and compare. Record experiments and report with the report template. This workflow grants no extra GPU, job stop, image publish or cross-repository write permission.
Numerical contracts and shared forward
Check the arithmetic that actually executes
The same formula does not guarantee the same floating-point result. Check accumulation dtype, cast position, FMA, reduction tree, chunking/split-K, approximate activations, quantization scale/rounding, epilogue and the input dtype of the final logprob. The same source can also produce different execution configs through autotune, compile specialization or a different GPU.
A shared kernel still needs its caller, arguments, layout and compile boundary checked; a direct call and a torch.compile fused path may differ. If needed, keep the same forward behind an existing opaque custom-op boundary; do not introduce unproven arithmetic rewrites to imitate the output.
Quantization execution checks (scale taken from the actual runner, quantization dispatch, quantizer outside the nn.Module tree, training vs deployment quantization grid) are in the precision, quantization and QAT guide, the grid in P8; format and scale records are in conversion chain §1.
First check that the semantics exist on both sides. When a feature is inside the kernel on one side and external or missing on the other, the difference does not go away with precision settings. Run a two-layer minimal comparison per model feature (attention sink, RoPE variant, final norm, etc.), enabling one at a time. In the external MiMo case SWA alone gave 0 and SWA+sink gave 3.906e-03, because only one side's kernel included the sink.
Discrete decisions come before last-bit error
After router/indexer logits, also compare the selected set, the order and the tie-break. Top-K with the same set in a different order still changes downstream gather/reduction. A deterministic tie-break defines the index order for equal scores and keeps the invalid/padding index contract.
For MoE, also fix or verify dispatch, expert capacity, grouped GEMM geometry, quantization and combine. Do not hide a router difference by forcing expert indices; forced routing only isolates the downstream.
The top-k/top-p truncation set is also a discrete decision: a last-bit logit difference moves tokens near the cutoff in or out of the set, and the normalization denominator of the processed logprob jumps with it. Compare the candidate set before processed logprob. MiMo-V2.6 replays the candidate set together with R3 in production: rollout records a fixed-shape, full-vocabulary bitmap (avoiding host sync and truncation), and training renormalizes within the set. If the contract adopts this replay, state in the report that it masks the set difference, and count the set flip rate separately.
Test shape and process boundaries
For the same victim, construct batch1 / multiple neighbors, different physical positions, lengths on both sides of a boundary, cross-process autotune, prefill/decode and Graph/eager. Different launches when only M changes are not a numerical difference; measure the output. Finding no difference covers only the shapes tested.
SSM chunk recurrence and per-token recurrence can be two legal but different rounding paths. Unifying them requires matching the chunk grid, boundary state precision and gate order, not just a deterministic switch.
For history-dependent operators, a single token's scan/conv input and output can make an upstream history difference look like an error in that operator. First find the first differing token at each layer over the full prompt batch, then capture norm, projection, conv, scan, gate and output projection in order at that layer. A single-position probe only narrows candidates; it does not prove the first cause. A plain Python hook in a compiled graph may hold stale warmup data; verify capture freshness with a predictable sentinel such as embedding/residual, and if needed use the custom-op boundary that actually executes and a separate buffer.
Boundary between shared forward and backward
With a shared inference forward + native differentiable VJP, verify separately: forward parity, a finite gradient for every trainable parameter, and gradient comparison under a cotangent representable in the same dtype. Operator-level VJP verification does not prove the whole-model gradient is bitwise equal to the original training implementation.
Do not detach the inference output and use it as a trainable substitute. A finite gradient does not mean effective learning; also check that advantage/reward is not degenerate, the gradient is non-zero, and parameters and optimizer state are updated.
Changing the arithmetic baseline
If you deliberately change both sides' arithmetic baseline to align them, keep the old baseline, build a new frozen rollout, state that "old outputs are not guaranteed to stay bitwise unchanged", and run model quality and performance regressions. Success numbers from the old baseline do not carry over to the new implementation.