dump-bisect-debug

## 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 1. **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. 2. **The reference must actually run end-to-end**, not "the code looks right" — references can have bugs too. 3. **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). 4. **Share as much code as possible**: reference and target use the same tokenizer / kernel / inputs, so the only variable is the code under investigation. 5. **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". 6. **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. 7. **Cross-check against a second reference**: besides the canonical reference, look at how sglang / vLLM implemented and fixed the same model. 8. **Dump names must be semantic**: `intra_attn_norm_in` / `intra_ffn_out`, never `tensor_5`. 9. **Dump one layer / one prompt at a time**: avoids file explosion and cross-contamination. 10. **Align `input_ids` first**: tokenizer mismatch is the most common false-positive bug source. 11. **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: 1. **The official repo's `inference/generate.py`**: run with `torchrun` (e.g. `/data/DeepSeek-V4-Pro/inference/generate.py`). 2. **HF transformers**: `AutoModelForCausalLM.from_pretrained(...).generate()`. 3. **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)`: ```bash 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): ```bash # ATOM side 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`): ```python 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(...) # writes ref_layer{LL}_Block__{stage}_rank{R}.pt to ATOM_FWD_DUMP_DIR ``` 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) ```bash # Reference 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 # Target (ATOM) ATOM_FWD_DUMP_DIR=$DIR ATOM_FWD_DUMP_LAYERS=0 \ python -m atom.examples.simple_inference --prompt "1+2+3=?" --max-tokens 1 # Compare 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):