source docs/recipes/balanced-training-mix.md · view on GitHub ↗

Balanced training mix

By default a trace-dump corpus (B22 → B29) is single-hop heavy — those are the queries that succeed soonest in real agent loops, so they accumulate fastest. Training on the raw blob bakes in that imbalance and makes the resulting specialist worse at multi-hop and comparison prompts even though both are present in the data.

This recipe runs B30 (posttrainllm reasoning-classify) over the synthesized JSONL before B29 hands it to posttrainllm sft, so the depth ratio is a knob the operator sets explicitly.

Pipeline

.atraj rollouts            (B22)

posttrainllm traces-to-data     (B29)  →  raw-sft.jsonl

posttrainllm reasoning-classify --score      → scored.jsonl
posttrainllm reasoning-classify --filter     → balanced.jsonl

posttrainllm sft --data balanced.jsonl       (Tier A)

Step 1 — train the depth classifier (one-time per dataset family)

A small hand-labeled seed (~500 prompts spanning the 4 classes) is enough to bootstrap. Castform’s published mix examples use a 4-of-{single-hop, multi-hop, comparison, other} split; the same labels apply here.

posttrainllm reasoning-classify \
  --train labeled-seed.jsonl \
  --heldout labeled-heldout.jsonl \
  --out reason.tgfr

Expected output: per-class precision / recall / F1, model written to reason.tgfr (~256 KB at default vocab 65 536).

Step 2 — score the trace-dump corpus

posttrainllm reasoning-classify \
  --score raw-sft.jsonl \
  --model reason.tgfr \
  --out scored.jsonl

Each row gains a reasoning_depth field. Throughput is bound by tokenization (~hundreds of MB/s on M-series), not arithmetic.

Step 3 — downsample to a target mix

posttrainllm reasoning-classify \
  --filter scored.jsonl \
  --target-mix "single=0.3,multi=0.5,comparison=0.2,other=0.0" \
  --out balanced.jsonl

Aliases (singlesingle-hop, multimulti-hop, comparecomparison) are accepted. The largest feasible N is chosen so the target ratios are achievable from on-disk supply; rows in the smallest class become the bottleneck.

A typical first cut: single=0.3,multi=0.5,comparison=0.2 — Castform-aligned, biased toward multi-hop to lift the depth that most-often regresses on raw trace dumps.

Step 4 — train

posttrainllm sft --data balanced.jsonl --base <gallery-pin> ...

Notes