| name | dump-bisect-debug |
| description | Locate forward numerical bugs by dumping intermediate tensors from a target implementation and a known-good reference, then bisecting layer by layer. Also covers batch-invariance bisect (the same token at any batch position should produce a bitwise-identical output, per DeepSeek V4 paper §3.3). Use when "the output is wrong but I don't know where" — model produces gibberish, degenerates, or picks the wrong token, but code review reveals nothing. |
| version | 3.1.0 |
| scope | ATOM (generalizable to any PyTorch forward debug) |
| last_updated | 2026-05-20T00:00:00.000Z |
ATOM built-in infrastructure (v3.0+ — read first)
Use these instead of hand-rolling hooks / compare scripts.
| Module | Purpose |
|---|
atom/utils/debug_helper/dump.py | Env-gated forward / weight / sampler dump; multi-class hooks; multi-call counter. |
atom/utils/debug_helper/compare.py | cos_max (double-precision — avoids the fp32 cos > 1.0 trap) + slot_split + pick_prefill_call + CLI. |
atom/utils/debug_helper/ref_patch.py | Monkey-patch context managers for instrumenting read-only reference implementations. |
atom/utils/envs.py "Debug Dump" section | 9 env vars (ATOM_FWD_DUMP_*, ATOM_WEIGHT_DUMP_*, ATOM_DEBUG_TOPK*), all no-op by default. |
Public API: from atom.utils.debug_helper import .... CLI: python -m atom.utils.debug_helper.compare <subcommand>.
Once these are wired up, debugging a new model never requires editing model_runner.py, writing a cos_max again, or re-inventing the warmup-vs-prefill heuristic.
Dump-Bisect Debug Methodology
Looking at logits or output text only tells you that it's wrong, not where it's wrong.
This methodology compresses where down to O(log N) dump-and-compare iterations.
Quick-start decision tree
Is the symptom reproducible at temperature=0.0 / single prompt?
├─ NO → batch / sampling-stochastic
│ └─ Same prompt yields different tokens across runs OR across batch slots?
│ ├─ across runs → you ALREADY BROKE determinism — fix sampling seed or env
│ └─ across slots → likely NOT a bug; jump to Phase 8 (batch invariance)
└─ YES → deterministic forward bug, run the linear pipeline:
├─ Have a reference impl? ──NO──► build one (Phase 0)
├─ Schema/shape mismatch? ──YES─► fix loader (Phase 1, do NOT skip)
├─ Single layer dump cos<1? ───── Phase 3 → Phase 4 → Phase 5
└─ Pinpoint sub-stage cos drop ── Phase 6 (standalone GPU isolation)
└─ Phase 7 (fix + isolation revert)
Phase-at-a-glance
| # | Goal | Time | Move on when |
|---|
| 0 | Establish reference | once per project | reference reproduces user-expected output end-to-end |
| 1 | Rule out weight loading | ≤ 1h, 1–2 GPU runs | all dumped params byte-equal or cos > 0.9999 |
| 2 | Define dump protocol | once per model | both sides agree on stage names + tensor contract |
| 3 | First single-layer comparison | ≤ 30min, 2 GPU runs | layer 0 cos table drawn |
| 4 | Layer-level bisect | 1–2h | first layer with cos < 0.99 (or rel > 10%) located |
| 5 | Intra-layer sub-stage bisect | 1–3h per bug | sub-stage with single-step cos drop > 0.001 located |
| 6 | Standalone GPU kernel isolation | 1–2h | path A reproduces ATOM dump cos > 0.9999; cross-experiment confirms quant vs GEMM |
| 7 | Fix + isolation revert | ~1h | e2e byte-equal vs ref AND each fix's necessity verified |
| 8 | Batch invariance bisect (parallel mode) | 2–3h | slot-vs-slot cos vs spec — if broken, classify as kernel-stack issue, NOT model bug |
Total per root cause: 4–6h. Multiple root causes stack. 5–10× faster than "stare at the code and guess" with no missed diagnoses.
When to use
Trigger conditions (any one applies):
- Model output is gibberish, degenerates, or picks the wrong token, but code review finds nothing.
- Output is correct on some prompts and wrong on others ("corner case").
- Single prompt OK / batch fails; prefill OK / decode fails; short prompt OK / long prompt fails.
- A correct reference implementation exists (HF transformers, official inference repo, a previously-verified commit).
Don't use when:
- No reference ground truth (only "I think it's wrong") — build a reference first.
- Difference of just a few tokens — could be numerical noise; confirm it's a real bug first.
- Symptom is "non-deterministic across runs at temp=0.0" — that's broken sampling determinism (env / RNG / kernel), not a forward bug; this skill won't help.
Core principles
- Read the model paper before assuming a bug. Some divergence is expected by design of the model architecture, not a bug to fix. V4 paper §3.3 explicitly assumes batch-invariant kernels; running on a non-batch-invariant inference stack will flip edge-confidence tokens and that is not an ATOM bug. Always check the model's reproducibility / determinism claims first — if the runtime stack doesn't meet them, classify as a kernel-stack limitation and document, don't bisect.
- The reference must actually run end-to-end, not "the code looks right" — references can have bugs too.
- Reference may be batch=1 only. Many official
inference/model.py files (V4, parts of DeepSeek family) hardcode max_batch_size=1. You cannot directly compare batch>1 against them. Either modify the ref to support batches OR design experiments that work within the bsz=1 constraint (e.g. multiple seeded runs).
- Share as much code as possible: reference and target use the same tokenizer / kernel / inputs, so the only variable is the code under investigation.
- Rule out weight loading first: confirm byte-equality before any forward bisect. Schema diff (names + shapes + dtypes) BEFORE numerical comparison — a missing param prefix masquerades as a "forward bug".
- Fix isolation revert is mandatory after multi-bug fixes. When 2+ fixes land together, revert each one in turn to identify which are critical (output-changing) vs fine-tuning (cosmetic / perf-cost). This decides PR split granularity — critical fixes merge fast, fine-tuning fixes can wait.
- Cross-check against a second reference: besides the canonical reference, look at how sglang / vLLM implemented and fixed the same model.
- Dump names must be semantic:
intra_attn_norm_in / intra_ffn_out, never tensor_5.
- Dump one layer / one prompt at a time: avoids file explosion and cross-contamination.
- Align
input_ids first: tokenizer mismatch is the most common false-positive bug source.
- Same stage name ≠ numerically equivalent: confirm both sides have applied the same number of ops at dump time. The most common trap: ATOM dumps pre-all-reduce, ref dumps post-all-reduce.
Eight-phase flow
Phases 0–7 are the linear pipeline for ref-vs-target bisect. Phase 8 is a parallel mode for batch-invariance investigation; trigger it independently when the symptom is "single prompt OK, batch fails".
Phase 0: Establish a reference
If no reference exists yet, build one. Priority:
- The official repo's
inference/generate.py: run with torchrun (e.g. /data/DeepSeek-V4-Pro/inference/generate.py).
- HF transformers:
AutoModelForCausalLM.from_pretrained(...).generate().
- A known-good prior commit:
git checkout <commit> and run.
Reference requirements:
- Same weights (stream from safetensors directly; do not convert).
- Same tokenizer / chat template.
- Same GPU kernel when investigating numerical drift — use the same aiter / cuBLAS kernels in the reference, otherwise you can't distinguish "algorithm bug" from "kernel numerical drift".
- End-to-end verified: the reference output must match user expectation (e.g.
"1+2+3=?" answers "6").
Output: ref_full_generate.py or similar — capable of producing a ground-truth token sequence.
Phase 1: Rule out weight loading (before any forward bisect)
Why first: if weights load wrong, every subsequent forward comparison will mis-attribute the cause to the forward path. 10 minutes of weight comparison saves a day of forward bisect.
Use maybe_dump_weights_and_exit(self.model) from atom.utils.debug_helper, already wired in model_runner.py. It dumps params + buffers per rank and sys.exit(0):
ATOM_WEIGHT_DUMP_DIR=/path/to/dump \
ATOM_WEIGHT_DUMP_LAYERS=0,2 \
python -m atom.examples.simple_inference --model ... -tp 8
Comparison checklist:
| Check type | What to do |
|---|
| Schema diff first | List ATOM-only / ref-only / shape mismatch / dtype diff before any numerical comparison. |
| FP8 weight | Compare byte-equality after aiter.ops.shuffle.shuffle_weight(ref_w, layout=(16,16)). |
| FP8 scale (e8m0) | Cast via aiter.utility.fp4_utils.e8m0_to_f32(ref_s) to fp32, then compare. |
| TP-replicated layer | Per-rank ATOM byte-equal vs ref. |
| TP-sharded layer | torch.cat(ref_rank0..7, dim=tp_dim) vs full ATOM (or per-rank ATOM vs ref slice + shuffle). |
| Norm weight (BF16 vs FP32) | Different dtype, equal value — cast then cos. |
| MoE expert weights | Volume is huge; skip initially — assume the expert loader is consistent with other weights. |
Conclusion patterns:
- ✓ All byte-equal / cos > 0.9999 → weight loading is OK; proceed to Phase 2.
- ✗ Any byte mismatch → fix the loader first (WeightsMapper / shuffle / TP shard / quant_type).
Phase 2: Define the forward dump protocol
Both sides agree on the same checkpoint names and tensor shape contract.
Minimum set (dump per layer):
| Stage | Meaning | Use |
|---|
embed.input_ids | Input token ids | Confirm tokenization consistency |
embed.embed_out | Embedding output | Confirm lookup consistency |
layer{L}.hidden_in | Hidden entering the layer | Confirm previous layer's output |
layer{L}.attn_norm_out | After attention norm | Isolate norm differences |
layer{L}.attn_out | Attention output | Whole attention block |
layer{L}.ffn_norm_out | After FFN norm | Isolate norm |
layer{L}.ffn_out | FFN output | Whole FFN block |
layer{L}.hidden_out | Hidden leaving the layer | Feeds the next layer's comparison |
embed.final_h | Pre-lm_head hidden | Accumulated diff vs lm_head amplification |
embed.final_logits | Final logits | Vocab-space difference |
Use the built-in dump infrastructure (do not write hooks by hand):
ATOM_FWD_DUMP_DIR=/path/to/dump \
ATOM_FWD_DUMP_LAYERS=0 \
ATOM_FWD_DUMP_BLOCK_CLASS=Block \
python -m atom.examples.simple_inference --prompt "1+2+3=?" --max-tokens 1
For the reference side (often a read-only /data/<model>/inference/model.py):
from atom.utils.debug_helper import patch_block_forward, patch_module_dump
with patch_block_forward(ref_Block, layer_attr="layer_id", side_prefix="ref"):
ref_model.forward(...)
For deeper sub-stages (RoPE, q_norm, sparse_attn output …), patch the relevant method directly with patch_method and insert named dump(stage, tensor) calls — copy the original forward body verbatim and insert the dumps in between to avoid losing side effects.
Phase 3: First comparison (single layer, single prompt)
ATOM_FWD_DUMP_DIR=$DIR ATOM_FWD_DUMP_LAYERS=0 \
torchrun --nproc-per-node=8 ref_full_generate.py --prompt "1+2+3=?" --max-new-tokens 1
ATOM_FWD_DUMP_DIR=$DIR ATOM_FWD_DUMP_LAYERS=0 \
python -m atom.examples.simple_inference --prompt "1+2+3=?" --max-tokens 1
python -m atom.utils.debug_helper.compare ref-vs-target --dir $DIR
The CLI uses the project's standard cos_max (double precision). It prints a per-stage table with severity flags and asserts input_ids match before doing anything else.
Severity thresholds (look at both cos and rel; rel is more sensitive):
| cos | rel | Meaning | Action |
|---|
> 0.9999 | < 1% | Bit-equal range | ✓ OK |
0.99 ~ 0.9999 | 1 ~ 10% | Numerical drift (kernel/dtype) | ? Flag — watch for accumulation |
0.9 ~ 0.99 | 10 ~ 30% | Mild algorithmic drift / partial heads wrong | ✗ Bisect to sub-stage |
< 0.9 | > 30% | Real bug | ✗ Locate immediately |
≈ 0 or negative | > 50% | Total scramble / sign flip | ✗ Usually weight loading / shuffle bug |
Important: when cos and rel disagree, trust rel:
- When hidden values span large ranges (e.g.
max_abs = 1e5), cos is dominated by a few outliers and may read 0.9998 even though rel = 57%.
- A 50%+ per-element error → after
lm_head amplification, logits are completely wrong.
Phase 4: Layer-level bisect — find the first cos drop, with layer-class awareness
The first layer with cos < 0.99 or rel > 10% = the layer where the bug lives.
Key observation: layer class. When dumping multiple layers (0 / N / 2N / 3N) to look at the decay curve, focus on which layer's cos suddenly drops and correlate with that layer's class:
| Model | Layer-class examples | Investigation direction |
|---|
| DeepSeek-V4 | First N layers use hash routing; rest use sqrtsoftplus routing | Did the non-hash path miss a fix? |
| DeepSeek-V4 | compress_ratio=4 (sparse) vs =128 (window) | Inside the sparse path? |
| Qwen3-Next | Hybrid attention vs full attention | Attention-type branch? |
| MTP | Base layer vs MTP block | MTP path needs an independent fix? |
V4 example: layer 0/2 hidden_out cos = 1.0 ✓, but layer 3 suddenly drops to cos = 0.98. Layers 0/1/2 are hash routing (layer_id < n_hash_layers); layer 3+ takes select_experts(sqrtsoftplus).
→ Compare the hash path vs the sqrtsoftplus path → discover the latter is missing * routed_scaling_factor.
Accumulated drift vs algorithmic bug:
- Each layer cos = 0.999 but compounded to layer 60 it's 0.94 → kernel drift.
- One layer suddenly drops from 0.999 to 0.5 → algorithmic bug.
- Some type of layer (e.g. layer 3, 5, 7, …) consistently has poor cos while others are fine → layer-class branch bug.
Phase 5: Intra-layer sub-stage bisect → component-level root cause
Once the layer is located, dump finer checkpoints inside it. Each arrow = one dump:
Attention:
x_in → wq_a → q_norm → qr → wq_b → q_pre_norm → q_post_norm → q_post_rope → ┐
├ → sparse_attn → o_pre_invrope → o_post_invrope → wo_a → wo_b
x_in → wkv → kv_pre_norm → kv_post_norm → kv_post_rope → kv_after_quant → ┘
FFN (MoE):
