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

scripts/compare_diloco.py

@@ -24,6 +24,7 @@
 import math
 import os
 import re
+import secrets
 import shutil
 import signal
 import socket
@@ -37,6 +38,7 @@
 REPO_ROOT = Path(__file__).resolve().parent.parent
 sys.path.insert(0, str(REPO_ROOT))
 
+from yeto.protocol import ISLAND_HMAC_KEY_ENV  # noqa: E402
 from yeto.benchmark_resume import (  # noqa: E402
     build_data_manifest,
     implementation_fingerprint,
@@ -1706,8 +1708,17 @@ def print_plan(args, arms: list[Arm]) -> None:
     print(f"  repetitions: {len(seeds)} training seed(s)")
 
 
+def ensure_island_hmac_key(environ=os.environ) -> None:
+    """The syncer now refuses unauthenticated islands; every arm runs on this
+    host, so one random per-run key, inherited by the syncer and all learner
+    processes, satisfies it without weakening the real launch path."""
+    if not environ.get(ISLAND_HMAC_KEY_ENV):
+        environ[ISLAND_HMAC_KEY_ENV] = secrets.token_hex(32)
+
+
 def main(argv=None) -> int:
     args = build_parser().parse_args(argv)
+    ensure_island_hmac_key()
     if args.eval_only:
         result = evaluate_loss(
             args.model,

tests/test_codex_traces.py

@@ -358,3 +358,45 @@ def test_convert_paths_reads_directories_and_writes_rows(tmp_path):
 
     assert len(rows) == 1
     assert rows[0]["messages"][1]["content"] == "answer"
+
+
+def test_codex_session_tool_calls_become_tool_messages():
+    row = row_from_session_events(
+        [
+            {"type": "event_msg", "payload": {"type": "user_message", "message": "make page"}},
+            {
+                "type": "response_item",
+                "payload": {
+                    "type": "function_call",
+                    "name": "exec_command",
+                    "arguments": '{"cmd":"ls"}',
+                    "call_id": "c1",
+                },
+            },
+            {
+                "type": "response_item",
+                "payload": {"type": "function_call_output", "call_id": "c1", "output": "a.txt"},
+            },
+            {
+                "type": "response_item",
+                "payload": {
+                    "type": "custom_tool_call",
+                    "name": "apply_patch",
+                    "input": "*** Begin Patch",
+                    "call_id": "c2",
+                },
+            },
+            {
+                "type": "response_item",
+                "payload": {"type": "custom_tool_call_output", "call_id": "c2", "output": "Success"},
+            },
+            {"type": "event_msg", "payload": {"type": "agent_message", "message": "done"}},
+        ]
+    )
+
+    msgs = row["messages"]
+    assert [m["role"] for m in msgs] == ["user", "assistant", "tool", "assistant", "tool", "assistant"]
+    assert msgs[1]["tool_calls"][0]["function"] == {"name": "exec_command", "arguments": '{"cmd":"ls"}'}
+    assert msgs[2] == {"role": "tool", "content": "a.txt", "tool_call_id": "c1"}
+    assert msgs[3]["tool_calls"][0]["function"] == {"name": "apply_patch", "arguments": "*** Begin Patch"}
+    assert msgs[4]["tool_call_id"] == "c2"

tests/test_lm_benchmark.py

@@ -515,3 +515,12 @@ def test_eval_fairness_requires_same_held_out_target_count():
     benchmark.validate_eval_token_match(records, {"trained_tokens": 1234})
     with pytest.raises(RuntimeError, match="held-out target-token mismatch"):
         benchmark.validate_eval_token_match(records, {"trained_tokens": 1233})
+
+
+def test_benchmark_sets_one_island_hmac_key_for_syncer_and_learners():
+    env = {}
+    benchmark.ensure_island_hmac_key(env)
+    key = env["YETO_ISLAND_HMAC_KEY"]
+    assert len(key) == 64
+    benchmark.ensure_island_hmac_key(env)
+    assert env["YETO_ISLAND_HMAC_KEY"] == key  # an operator-provided key is kept

tests/test_qwen_generation_template.py

@@ -0,0 +1,95 @@
+"""Qwen3 / Qwen3.5 official templates gain exact native assistant masks.
+
+Uses tokenizers from the local Hugging Face cache only; each case skips when
+its tokenizer is not cached, so the test never needs the network.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from yeto.data import (
+    ExactAssistantMaskError,
+    _row_tokens,
+    normalize_sft_row,
+    qwen_generation_template,
+)
+
+TOOLS = [
+    {
+        "type": "function",
+        "function": {
+            "name": "f",
+            "description": "d",
+            "parameters": {"type": "object", "properties": {"x": {"type": "integer"}}},
+        },
+    }
+]
+
+ROW = {
+    "messages": [
+        {"role": "system", "content": "SYSMARK"},
+        {"role": "user", "content": "USERMARK one"},
+        {"role": "assistant", "content": "hist", "reasoning_content": "old thought"},
+        {"role": "user", "content": "USERMARK two"},
+        {
+            "role": "assistant",
+            "content": "<think>\nnew thought\n</think>\n\ncalling",
+            "tool_calls": [
+                {"type": "function", "function": {"name": "f", "arguments": {"x": 1}}}
+            ],
+        },
+        {"role": "tool", "content": "TOOLMARK"},
+        {"role": "assistant", "content": "", "reasoning_content": "r3"},
+        {"role": "assistant", "content": "final"},
+    ]
+}
+
+
+def _cached_tokenizer(name):
+    transformers = pytest.importorskip("transformers")
+    try:
+        return transformers.AutoTokenizer.from_pretrained(name, local_files_only=True)
+    except Exception as exc:  # not in the local cache
+        pytest.skip(f"{name} tokenizer not cached: {exc}")
+
+
+@pytest.mark.parametrize("name", ["Qwen/Qwen3-0.6B", "Qwen/Qwen3.5-0.8B"])
+@pytest.mark.parametrize("tools", [None, TOOLS])
+def test_qwen_native_mask_matches_official_ids_and_covers_only_assistant(name, tools):
+    tokenizer = _cached_tokenizer(name)
+    row = dict(ROW, tools=tools) if tools else ROW
+    example = normalize_sft_row(row, "auto", row_index=0)
+    official = tokenizer.apply_chat_template(
+        example.messages, tools=example.tools, tokenize=True, return_dict=True
+    )["input_ids"]
+
+    ids, weights = _row_tokens(tokenizer, row, assistant_mask_mode="native")
+
+    assert ids == list(official)
+    trained = tokenizer.decode([t for t, w in zip(ids, weights) if w])
+    untrained = tokenizer.decode([t for t, w in zip(ids, weights) if not w])
+    for marker in ("SYSMARK", "USERMARK", "TOOLMARK", "<tool_response>", "system", "old thought"):
+        assert marker not in trained
+    assert "<|im_start|>" not in trained
+    for marker in ("hist", "new thought", "calling", "<tool_call>", "r3", "final"):
+        assert marker in trained
+    assert trained.count("<|im_end|>") == 4
+    assert "TOOLMARK" in untrained and "USERMARK two" in untrained
+
+
+def test_rewrite_leaves_non_qwen_templates_alone():
+    plain = "{% for message in messages %}{{ message.content }}{% endfor %}"
+    assert qwen_generation_template(plain) is None
+
+
+def test_non_qwen_template_without_generation_still_fails():
+    class Tok:
+        name_or_path = "other"
+        chat_template = "{% for message in messages %}{{ message.content }}{% endfor %}"
+
+        def get_chat_template(self, chat_template=None, tools=None):
+            return self.chat_template
+
+    with pytest.raises(ExactAssistantMaskError, match="does not contain"):
+        _row_tokens(Tok(), {"messages": [{"role": "user", "content": "u"}, {"role": "assistant", "content": "a"}]})

tools/eval_sft_heldout.py

@@ -0,0 +1,113 @@
+"""Weighted next-token CE of a base model (+ optional PEFT adapter) on SFT rows.
+
+Blocks are built with yeto.data._target_packed_blocks (row file order, no
+shuffle), i.e. the same _row_tokens + _append_target_blocks path the learner
+uses. Loss/token = sum(w[t+1] * CE(t -> t+1)) / sum(w[t+1]) over all blocks,
+matching the learner's one-token shift (yeto/learner.py target count).
+
+Usage:
+  PYTHONPATH=. python tools/eval_sft_heldout.py --model Qwen/Qwen3-0.6B \
+      [--model-revision REV] [--adapter DIR] --data eval.jsonl \
+      [--seq-len 4096] [--assistant-mask-mode legacy|native] [--max-blocks N]
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+
+import torch
+import torch.nn.functional as F
+
+from yeto.data import _target_packed_blocks, load_rows
+
+
+def build_blocks(tokenizer, data, seq_len, mask_mode, data_format="auto", max_rows=None):
+    rows = load_rows(data)
+    n = len(rows) if max_rows is None else min(len(rows), max_rows)
+    return _target_packed_blocks(
+        rows, tokenizer, list(range(n)), seq_len, "assistant", mask_mode, data_format
+    )
+
+
+def load_model(model_id, revision, adapter, dtype):
+    from transformers import AutoModelForCausalLM
+
+    model = AutoModelForCausalLM.from_pretrained(model_id, revision=revision, dtype=dtype)
+    info = None
+    if adapter:
+        from peft import PeftModel
+        from peft.utils import set_peft_model_state_dict
+        from safetensors.torch import load_file
+
+        model = PeftModel.from_pretrained(model, adapter)
+        # Re-apply explicitly to surface any key mismatch as an error.
+        sd = load_file(f"{adapter}/adapter_model.safetensors")
+        res = set_peft_model_state_dict(model, sd)
+        info = {
+            "adapter_tensors": len(sd),
+            "missing_lora_keys": [k for k in res.missing_keys if "lora_" in k],
+            "unexpected_keys": list(res.unexpected_keys),
+        }
+        if info["missing_lora_keys"] or info["unexpected_keys"]:
+            raise SystemExit(f"adapter key mismatch: {json.dumps(info)[:2000]}")
+    model.eval()
+    return model, info
+
+
+@torch.no_grad()
+def weighted_ce(model, blocks, weights):
+    total_loss, total_w = 0.0, 0.0
+    for ids, w in zip(blocks, weights):
+        x = torch.tensor([ids], dtype=torch.long)
+        wt = torch.tensor(w[1:], dtype=torch.float64)
+        logits = model(input_ids=x).logits[0, :-1].float()
+        ce = F.cross_entropy(logits, x[0, 1:], reduction="none").double()
+        total_loss += float((ce * wt).sum())
+        total_w += float(wt.sum())
+    return total_loss, total_w
+
+
+def main(argv=None):
+    p = argparse.ArgumentParser()
+    p.add_argument("--model", required=True)
+    p.add_argument("--model-revision", default=None)
+    p.add_argument("--adapter", default=None)
+    p.add_argument("--data", required=True)
+    p.add_argument("--data-format", default="auto")
+    p.add_argument("--seq-len", type=int, default=4096)
+    p.add_argument("--assistant-mask-mode", choices=["legacy", "native"], default="legacy")
+    p.add_argument("--max-blocks", type=int, default=None)
+    p.add_argument("--max-rows", type=int, default=None)
+    p.add_argument("--dtype", choices=["float32", "bfloat16"], default="float32")
+    p.add_argument("--threads", type=int, default=None)
+    p.add_argument("--output", default=None)
+    a = p.parse_args(argv)
+    if a.threads:
+        torch.set_num_threads(a.threads)
+    from transformers import AutoTokenizer
+
+    tok = AutoTokenizer.from_pretrained(a.adapter or a.model, revision=None if a.adapter else a.model_revision)
+    blocks, weights = build_blocks(tok, a.data, a.seq_len, a.assistant_mask_mode, a.data_format, a.max_rows)
+    total_blocks = len(blocks)
+    if a.max_blocks is not None:
+        blocks, weights = blocks[: a.max_blocks], weights[: a.max_blocks]
+    model, info = load_model(a.model, a.model_revision, a.adapter, getattr(torch, a.dtype))
+    loss, tw = weighted_ce(model, blocks, weights)
+    out = {
+        "model": a.model, "model_revision": a.model_revision, "adapter": a.adapter,
+        "data": a.data, "seq_len": a.seq_len, "assistant_mask_mode": a.assistant_mask_mode,
+        "dtype": a.dtype, "blocks_total": total_blocks, "blocks_used": len(blocks),
+        "total_tokens": sum(len(b) for b in blocks),
+        "target_tokens": tw, "loss_sum": loss,
+        "loss_per_token": loss / tw if tw else None,
+        "adapter_check": info,
+    }
+    s = json.dumps(out, indent=2)
+    if a.output:
+        open(a.output, "w").write(s + "\n")
+    print(s)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tools/test_eval_sft_heldout.py

@@ -0,0 +1,47 @@
+"""Self-checks for tools/eval_sft_heldout.py (no network, tiny fake model).
+
+Run: PYTHONPATH=. python -m pytest tools/test_eval_sft_heldout.py
+Optional real check: set RECOVERED_ADAPTER=<dir> to verify a recovered PEFT
+adapter loads onto Qwen/Qwen3-0.6B with no missing/unexpected keys.
+"""
+import os
+import sys
+
+import pytest
+import torch
+
+sys.path.insert(0, os.path.dirname(__file__))
+from eval_sft_heldout import weighted_ce  # noqa: E402
+
+
+class _Fixed(torch.nn.Module):
+    def __init__(self, vocab):
+        super().__init__()
+        self.vocab = vocab
+
+    def forward(self, input_ids):
+        from types import SimpleNamespace
+
+        return SimpleNamespace(logits=torch.zeros(*input_ids.shape, self.vocab))
+
+
+def test_weighted_ce_uses_shifted_weights():
+    vocab = 7
+    loss, tw = weighted_ce(_Fixed(vocab), [[1, 2, 3, 4]], [[1.0, 0.0, 1.0, 1.0]])
+    # weights[1:] = [0, 1, 1] -> 2 targets, uniform logits -> CE = log(vocab)
+    assert tw == 2.0
+    assert loss / tw == pytest.approx(torch.log(torch.tensor(float(vocab))).item())
+
+
+@pytest.mark.skipif(not os.environ.get("RECOVERED_ADAPTER"), reason="set RECOVERED_ADAPTER")
+def test_recovered_adapter_loads_cleanly():
+    from eval_sft_heldout import load_model
+
+    _, info = load_model(
+        "Qwen/Qwen3-0.6B",
+        "c1899de289a04d12100db370d81485cdf75e47ca",
+        os.environ["RECOVERED_ADAPTER"],
+        torch.float32,
+    )
+    assert info["adapter_tensors"] == 392
+    assert not info["missing_lora_keys"] and not info["unexpected_keys"]

yeto/codex_traces.py

@@ -480,6 +480,29 @@ def row_from_session_events(
             if payload_type == "reasoning" and _looks_encrypted_reasoning(payload):
                 skipped_encrypted_reasoning += 1
                 pending_encrypted_reasoning += 1
+            elif payload_type in {"function_call", "custom_tool_call", "local_shell_call"}:
+                # Codex rollouts record tool use only as response items; keep
+                # them so agentic rows carry assistant tool_calls + tool results.
+                call_id = payload.get("call_id") or payload.get("id")
+                tool_call = _tool_call(
+                    {
+                        "tool_call_id": call_id,
+                        "name": payload.get("name") or payload_type,
+                        "arguments": payload.get("arguments", payload.get("input", payload.get("action", {}))),
+                    },
+                    len(messages),
+                )
+                _append_message(messages, "assistant", "", tool_calls=[tool_call])
+            elif payload_type in {"function_call_output", "custom_tool_call_output", "local_shell_call_output"}:
+                output = payload.get("output")
+                if isinstance(output, dict):
+                    output = output.get("content", output)
+                _append_message(
+                    messages,
+                    "tool",
+                    _truncate(_codex_output_text(output), max_tool_output_chars),
+                    tool_call_id=payload.get("call_id"),
+                )
             elif payload_type == "message" and not has_codex_event_messages:
                 role = payload.get("role")
                 content = _codex_output_text(payload.get("content"))