File size: 5,118 Bytes
85b17bd f6158c7 85b17bd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | # DAPO Math Jacobian matrices
Fits averaged Jacobian transport matrices from Anthropic's Jacobian Lens
method on a local Qwen checkpoint and DAPO Math prompts.
For block output `l` and target block output `L`:
```text
J_l = E_prompt,source-position,future-target-position[d h_L / d h_l]
```
The estimator follows Anthropic's Apache-2.0
[`jacobian-lens`](https://github.com/anthropics/jacobian-lens) reference code.
## Layer convention
Qwen3-4B has 36 transformer blocks indexed `0..35`. The default target is the
output of block 35, and matrices are fitted for block outputs `0..34`. The
target-to-itself map is exactly identity and is included as matrix 35, giving
one exported matrix for every transformer-block output. The embedding boundary
is not included in this first version.
## Install and smoke fit
```bash
python -m pip install -e .
fit-jacobians \
--model LLMs/qwen3-4b-base-sft-qwen3-8b \
--data data/dapo-math-17k/dapo-math-17k.jsonl \
--num-prompts 20 \
--dim-batch 8 \
--output-dir outputs/smoke-20
```
The resumable checkpoint keeps FP32 sums. `lens-bf16.pt` contains the final
BF16 means. Increase `--dim-batch` only if GPU memory permits.
For the main fit, change `--num-prompts` to `1000` and preferably set
`--checkpoint-every 10`.
## Multiple GPUs
The multi-GPU launcher fits disjoint prompt shards with one complete model
replica per GPU, then merges the FP32 sums exactly before converting to BF16:
```bash
GPUS=0,1,2,3 \
NUM_PROMPTS=1000 \
OUTPUT_DIR=outputs/main-1000 \
bash scripts/fit_multi_gpu.sh
```
For two GPUs and the default 20-prompt smoke test:
```bash
GPUS=0,1 bash scripts/fit_multi_gpu.sh
```
Optional variables include `DIM_BATCH`, `MAX_SEQ_LEN`, `CHECKPOINT_EVERY`,
`MODEL_PATH`, `DATA_PATH`, `SEED`, and `PYTHON_BIN`. Each process sees its
assigned physical GPU as `cuda:0`. Failed shards retain their checkpoints, so
running the same command resumes them rather than restarting.
JSONL input uses question text. Parquet input uses pre-tokenized student
response windows as described below; both share the same Jacobian estimator.
## Fit on Parquet student responses
Parquet input is auto-detected and reads `metadata.response_tokens` directly.
Selected rollouts rotate through four evenly spaced 1,024-token regions. For a
4,096-token rollout, their starts are 0, 1,024, 2,048, and 3,072. Labels,
response text, loss masks, and teacher log probabilities are not used.
Start with a 20-rollout multi-GPU smoke fit:
```bash
GPUS=0,1,2,3 \
NUM_PROMPTS=20 \
DATA_PATH=rollouts/dapo-math-17k-qwen3-4b-sft-rollouts-lightning-opd-precomputed.parquet \
CORPUS_FORMAT=rollout-parquet \
RESPONSE_WINDOW_LEN=1024 \
DIM_BATCH=8 \
CHECKPOINT_EVERY=1 \
OUTPUT_DIR=outputs/response-1024-smoke-20 \
bash scripts/fit_multi_gpu.sh
```
When `MAX_SEQ_LEN` is omitted, Parquet runs automatically use
`RESPONSE_WINDOW_LEN`; JSONL runs use 128. A 1,024-token sequence is much more
expensive than a 128-token sequence, so begin with `DIM_BATCH=8` and increase
to 16 or 32 only after checking H100 memory and time per rollout.
## Interactive sanity check
Use held-out samples after the fitting range. For a lens fitted with the first
20 seed-17 shuffled prompts:
```bash
python -m math_jlens.explore \
--model LLMs/qwen3-4b-base-sft-qwen3-8b \
--data data/dapo-math-17k/dapo-math-17k.jsonl \
--lens outputs/multi-gpu-smoke-20/lens-bf16.pt \
--fit-count 20 \
--seed 17 \
--sample-index 0 \
--device cuda:0
```
The explorer generates one deterministic solution, caches it, performs a
teacher-forced residual capture, and opens a prompt. Useful commands:
```text
tokens [start] [end] show token positions (P=prompt, R=response)
inspect POSITION [TOP_K] top J-lens tokens across every layer
compare POSITION LAYER compare J-lens against ordinary logit lens
trace "TOKEN TEXT" [POSITION] show token ranks across layers
save START END [LAYER|all] [TOP_K] [FILE]
export original and J-lens tokens to text
answer trace the first gold-answer token before output
problem show the complete DAPO problem
output show the complete generated solution
result show extracted answer, ground truth, and match
info show sample, problem, output, and result together
sample INDEX load another held-out sample
next load the next held-out sample
quit
```
At activation position `t`, the readout predicts token `t+1`; the interface
always prints both tokens to prevent an off-by-one interpretation.
Examples of text export:
```text
save 200 220 all layers, top 5, automatic filename
save 200 220 18 layer 18 only, top 5
save 200 220 all 10 all layers, top 10
save 200 220 18 10 report.txt layer 18, top 10, explicit filename
```
`END` is exclusive. Automatic files are written under
`outputs/jlens-exports/` and include the original token, next predicted token,
token IDs, mapped J-lens tokens, ranks, and logits.
|