Codeseys's picture
examples: add sdpo_real_trace_train_smoke — close the forward+backward+step link
7090729
|
Raw
History Blame Contribute Delete
4.13 kB

SDPO real-trace training smoke

The missing forward + backward + optimizer step link for SDPO.

Why this exists

The framework proved two halves of the SDPO loop in isolation:

Half Where What it proves
Data path examples/validate_real_trace_alignment/ ingestion → adapter → collator emits a batch whose sdpo_loss_mask lands on content tokens at 100% alignment on real `/.claude` traces
Loss math composer_replication/tests/test_gradient_flow.py compose_loss routes finite non-zero gradients through the SDPO channel — but only on a millisecond TinyLM stand-in (no HF model)

Nobody had connected them: an actual compose_loss forward + backward + optimizer.step() on a real HuggingFace model fed by the real-trace collator. That is the one unproven edge — and it is exactly the never-implemented composer_replication.examples.sdpo_with_real_traces_production module that the Modal stage_4_sdpo_smoke referenced. This script is that module, made real.

What it asserts (the gates)

  1. The collated real-trace batch drives compose_loss without crashing.
  2. total loss is finite (not NaN/Inf) across all steps.
  3. The SDPO channel fires: sdpo_jsd > 0 on ≥1 step — proves the shape-gate at loss.py:163 passed and the hint-conditioned teacher forward contributed real signal (not the silent no-op the empty-placeholder stage_4 would give).
  4. A real parameter moved after optimizer.step() (training happened).

Run

# Canonical PASS config (B=1 fp32 — fast native CPU GEMM, ~14GB peak):
python examples/sdpo_real_trace_train_smoke/run.py \
    --max-sessions 6 --max-steps 2 --max-examples 1 --dtype fp32

Verified PASS (Qwen2.5-0.5B-Instruct, CPU, 6 real ~/.claude error sessions):

collated batch: input_ids (1, 1339), sdpo_loss_mask in-loss positions = 6
  step 0: total=2.36307  lm_ce=2.33588  sdpo_jsd=0.02718  finite=True
  step 1: total=2.32758  lm_ce=2.30190  sdpo_jsd=0.02568  finite=True
  all losses finite:        True
  SDPO channel fired (>0):  True
  param 'model.embed_tokens.weight' moved: True  (max|Δ|=6.22e-05)
  RESULT: PASS ✅

Operational notes (hard-won)

  • Target model = small instruct (Qwen2.5-0.5B-Instruct), NOT nanochat. Agent-trace SDPO needs traces with tool-error → recovery structure. A trained nanochat is a plain chat model with no tool-use → 0% SDPO error sites by construction. The correct SDPO target is a small instruct model with a chat template.
  • Memory: the killer is vocab × seq × dtype. Qwen2.5 vocab is 151,936, so fp32 logits are ~1.17 GB per (example, 2048-tok) forward; SDPO does two forwards (student + hint-conditioned teacher). The fp32 forward+backward transiently hits ~27 GB and trips the host/cgroup OOM killer at B≥2. B=1 fp32 keeps the peak ~14 GB and uses fast native CPU GEMM.
  • Do NOT use bf16 on CPU for this. bf16 clears the memory wall but CPUs without AVX512-BF16 fall back to emulated GEMM — a >10× slowdown (a single step ran >13 min vs ~30-60 s in fp32). The --dtype bf16 flag exists but fp32 + B=1 is the fast path.
  • Sequence length carries the signal — do not over-truncate. The error-recovery turns sit deep in long agent sessions. --max-seq-len 1024 truncated all SDPO sites away → all-zero mask → SKIP. Keep ≥1536; the script SKIP-guards (exit 2) rather than silently training on zero signal.
  • --strip-thinking defaults False (correct for SDPO): on real Claude Code traces the recovery turn is frequently pure [THINKING]; stripping empties ~67% of error sites and the SDPO channel sees no signal.
  • Run it detached from the gateway cgroup if iterating live: systemd-run --user --scope -p MemoryMax=28G -- .... A gateway restart SIGTERMs every child in its cgroup (exit 143); a transient scope survives.

Exit codes

  • 0 PASS (all gates)
  • 1 FAIL (a gate failed — non-finite loss, SDPO never fired, or no param moved)
  • 2 SKIP (no error-bearing sessions, no chat-template model, or mask all-zero)