Agentic SFT on Codex traces: keep tool calls, native Qwen mask, held-out eval, benchmark HMAC fix

[michaellchung]

Summary

Make multi-island (Decoupled DiLoCo) agentic SFT on Codex traces work end to end, and fix two things that broke on current main.

1. codex_traces: keep tool calls. Codex rollouts record tool use only as response_item function_call / custom_tool_call / local_shell_call (+ *_output). The converter dropped them, so "agentic" rows had no tool calls. One real rollout: 21 messages / 0 tool calls before → 161 messages / 70 tool calls after.
2. Exact native assistant mask for Qwen3 / Qwen3.5. The official templates have no {% generation %}, so --assistant-mask-mode native refused them and users fell back to the synthetic legacy format. When a template matches the Qwen shape, data.py now rewrites only the assistant branch to wrap the assistant output in {% generation %}. The rendered text is unchanged. Other templates behave as before: if the rewrite does not apply, the error is raised exactly as before.
   • Qwen3-0.6B, 20 held-out Codex rows: input_ids match the official apply_chat_template on 20/20. Every user/tool token is mask 0. The trained token share is 66.8%.
3. tools/eval_sft_heldout.py: held-out assistant-token CE for a base model plus an optional adapter. It packs exactly like training and checks that adapter keys match on load.
4. compare_diloco: per-run island HMAC key. The legacy syncer now refuses unauthenticated islands, so every DiLoCo arm of the LM benchmark died at syncer start. The harness now sets one random key per run, and the syncer and all learners inherit it. A key the operator already set is kept.

Evidence (real GPU runs, Modal L40S)

Data: HF AletheiaResearch/GPT-5.3-Codex-Spark-Codex, 200 real Codex CLI rollouts. Each rollout's originator is codex_exec and its model is gpt-5.3-codex-spark. Reasoning is encrypted and was skipped. The dataset card states no license; we used it for internal testing only. Model: Qwen3-0.6B with LoRA r16 on all linear layers, seq 4096. Loss counts assistant tokens only.

Fair same-host comparison (scripts/compare_diloco.py, 1 container × 2 L40S, 2 seeds, 4.19M raw tokens per arm, 20 held-out rows):

arm	held-out CE/token	Δ vs sync
base (untrained)	1.1300	–
baseline-m2 (one 2-rank sync group)	0.8236 ± 0.0000	–
m2 (2 DiLoCo islands × 1 GPU)	0.8304 ± 0.0014	+0.83% ± 0.17

Both arms processed the same number of target tokens (2,462,499).

Real cross-container 2-island runs (yeto launch, 2 × Modal L40S:1, syncer on a Nebius CPU head):

• Run a used the default grace. The adaptive grace computed to 0 because τ·ξ_step ≈ 4.8 s while one round needs ≈ 10 s to reach quorum. The syncer rejected the later island's push 23 times (OutOfRound), and only 25/48 outer steps merged both islands.
• Run b used --grace-tau 8 --grace-ms 3000. The syncer rejected 4 pushes, and 44/48 outer steps merged both islands.
• The final adapter was rebuilt from the syncer checkpoint with yeto.export and passed the layout check.

Tests

tests/test_codex_traces.py tests/test_data.py tests/test_data_formats.py tests/test_qwen_generation_template.py tests/test_lm_benchmark.py tests/test_losses.py tests/test_launch_auto.py tests/test_head_mode.py tests/test_benchmark_resume.py tools/test_eval_sft_heldout.py → 167 passed, 1 skipped (offline HF cache).

Known issues (not fixed here)

• The default grace (--grace-tau 2) collapses to 0 for short inner steps with fragments=4. Islands with even small arrival skew then take turns instead of merging. This PR does not change that default. The numbers are above, for a follow-up decision.
• The head job exits with rc=2 (FAILED) after printing "run finished", even though both islands SUCCEEDED and the checkpoint is intact. This happened in both runs a and b.
• Qwen3.5 templates need tool-call arguments as a dict. Codex traces carry JSON strings, so 19 of the 20 held-out rows fail even with the official 3.5 template. The shared normalizer is left unchanged.
• Long rollouts (median ~40k tokens) are cut into seq-len windows, so later windows lose earlier context.

🤖 Generated with Claude Code

[michaellchung]

Held-out eval of real cross-container run b (adapter rebuilt from the syncer checkpoint with yeto.export, same 20 held-out rows, native mask, CPU, all 205 blocks):

model	CE/token
base Qwen3-0.6B	1.1294
run b (2 Modal islands, --grace-tau 8)	0.7510

The CPU base value matches the GPU harness base (1.1300). Run b processed ~59M raw tokens across both islands, versus 4.19M per arm in the benchmark table, so run b is not comparable with that table. This result only shows that the real multi-island path trains and generalizes to rows it never saw.