ayh015/J-space

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:

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

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:

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:

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:

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:

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:

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:

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.