README.md

# 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.