modular__skills/import-model/references/layer-by-layer-debugging.md

Reading the layer-by-layer debug output

For parity/coherence failures (server runs but output is wrong), load the debug-model skill first. That protocol mandates per-layer tensor-dump comparators and parallel investigation agents. This reference covers the quick compare_layers.py probe that import-model runs before handing off.

The divergence hunt has two layers of debugging:

1. compare_layers.py: automated logit probes (what the script actually runs).
2. Per-layer tensor dumps: HF vs MAX comparators (see debug-model; preferred over scalar ops.output() taps).

Running compare_layers.py

From the repo root (via the OSS shim):

pixi run python compare_layers.py <HF_MODEL_ID> \
  --slug <your_slug> \
  --port 8000 \
  --prompt "The capital of France is"

Requires pixi run max serve with --custom-architectures <port_dir> (the slug folder, not its parent).

Flags:

Flag	Purpose
--dtype bfloat16|float32	Must match --quantization-encoding on serve (default bfloat16)
--skip-hf-layers	Skip HF hidden-state stats (faster)
--skip-multi	Skip multi-position prefix probe
--long-prompt "..."	Prompt for multi-position probe (~110 tokens default)
--positions 1,5,20,50,100,200	Prefix lengths to compare

Preflight before serve (guard + scaffold):

pixi run python check_walls.py <HF_MODEL_ID>
pixi run python run_oss_gates.py <HF_MODEL_ID> --port-dir <port_dir>/

After serve (during the divergence hunt):

pixi run python run_oss_gates.py <HF_MODEL_ID> \
  --port-dir <port_dir>/ --phase verify --slug <slug> --port 8000

---

What the script prints

1. HF hidden-state stats (HF-only)

  layer        mean     max_abs        norm
    0      0.0001      0.4200       12.34
    ...

These rows are diagnostic only — MAX does not expose hidden states through the completions API. Use them to confirm HF loads sanely (no NaNs, reasonable norms). They do not compare against MAX.

2. Short-prompt top-1 logprob

       check            hf           max    rel_diff  verdict
top1_logprob      -0.1234      -0.1250      0.0130  ok

Verdict	Meaning
ok	Relative logprob difference < 5%
DIVERGED	Prefill logits disagree at the last prompt token

If this row diverges at token 0, the bug is early: tokenizer, embedding, norm order, or attention block 0.

3. Multi-position top-1 probe

   pos     hf_id              max_text  verdict
     1      1234                  ' the'  ok
     5      5678                  ' fox'  ok
    50      9012                  ' lap'  DIVERGED

The script decodes HF's argmax token id and compares to MAX's greedy completion text at each prefix length.

First DIVERGED row localizes position-dependent bugs:

First divergence at	Likely cause
Position 1–5	Tokenizer, embedding, early attention
Position 5–20	GQA repeat, QK-norm, head layout
Position 50+	RoPE style/theta, partial-RoPE padding
Position 100+	Decode-path position handling, sliding window

See divergences.md for fixes.

---

Per-layer tensor dumps (when logit probes aren't enough)

When short and multi-position probes pass but output is still wrong, or you need to know which sublayer diverged, follow debug-model: build HF and MAX dumpers, run the comparator, and read comparator-output-patterns.md.

Use ops.output(...) taps only as fine-grained sub-taps after the lead agent localizes the broken subsystem, not as the primary bisect loop.

Interpreting tap diffs

If the first divergence is at:

• post-embed — check weight_adapters.py for embed_tokens.weight; check tie_word_embeddings.
• post-attn (layer 0) — attention: Q/K/V naming, RoPE style, GQA, QK-norm, sliding-window mask.
• post-mlp (layer 0) — MLP: activation variant (gelu_new vs gelu_tanh), gated layout, bias terms.
• post-block-norm but attn and MLP ok — block wiring: pre-norm vs parallel vs dual-norm, residual order, MuP multipliers.

Scaling vs direction bugs

• Similar direction, different magnitude — missing scale factor, wrong softmax scale, wrong norm formulation (x * weight vs x * (1 + weight)).
• Similar magnitude, low cos_sim — RoPE style mismatch, head permutation, sign error in rotate_half.

---

Iteration loop

1. Run compare_layers.py (or run_oss_gates.py --phase verify).
2. Localize from the first DIVERGED row (logprob row or multi-position row).
3. Open the matching HF forward() for that layer/component.
4. Fix one operation in <slug>.py or weight_adapters.py.
5. Restart pixi run max serve — stale compiled graphs hide fixes.
6. Re-run.

Each pass should push the first divergence later. If it doesn't, revert and re-read HF source.

---

Limitations

• MAX completions API has no per-layer hidden-state export — taps are manual.
• Top-1 text comparison assumes greedy decode; use temperature=0.0 on serve.
• HF reference loads with trust_remote_code=True and device_map="auto".
• Both sides must use the same dtype (--dtype must match serve encoding).

For trust_remote_code models that NaN with .to("cuda"), see divergences.md #19 and pitfalls-serving.md.