diff --git a/scripts/compare_diloco.py b/scripts/compare_diloco.py
index 1b3e6b9fd..22dc3450d 100644
--- a/scripts/compare_diloco.py
+++ b/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,
diff --git a/tests/test_codex_traces.py b/tests/test_codex_traces.py
index fb33a65fc..55bad9b7e 100644
--- a/tests/test_codex_traces.py
+++ b/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"
diff --git a/tests/test_lm_benchmark.py b/tests/test_lm_benchmark.py
index 6ed762e8f..3e0c5f3c6 100644
--- a/tests/test_lm_benchmark.py
+++ b/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
diff --git a/tests/test_qwen_generation_template.py b/tests/test_qwen_generation_template.py
new file mode 100644
index 000000000..0ca7fa3a5
--- /dev/null
+++ b/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": "\nnew thought\n\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", "", "system", "old thought"):
+ assert marker not in trained
+ assert "<|im_start|>" not in trained
+ for marker in ("hist", "new thought", "calling", "", "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"}]})
diff --git a/tools/eval_sft_heldout.py b/tools/eval_sft_heldout.py
new file mode 100644
index 000000000..19fbf2184
--- /dev/null
+++ b/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())
diff --git a/tools/test_eval_sft_heldout.py b/tools/test_eval_sft_heldout.py
new file mode 100644
index 000000000..74dbf73b3
--- /dev/null
+++ b/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= 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"]
diff --git a/yeto/codex_traces.py b/yeto/codex_traces.py
index 1e7ab4bb6..2ee283250 100644
--- a/yeto/codex_traces.py
+++ b/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"))
diff --git a/yeto/data.py b/yeto/data.py
index 3e7c986f8..f99e0252f 100644
--- a/yeto/data.py
+++ b/yeto/data.py
@@ -188,6 +188,19 @@ def _legacy_mask_hint() -> str:
)
+def _plain_selected_template(tokenizer, tools: list | None):
+ getter = getattr(tokenizer, "get_chat_template", None)
+ try:
+ if callable(getter):
+ return getter(chat_template=None, tools=tools)
+ return getattr(tokenizer, "chat_template", None)
+ except Exception as exc:
+ raise ExactAssistantMaskError(
+ f"tokenizer {_tokenizer_label(tokenizer)!r} could not select its native "
+ f"chat template: {exc}. " + _legacy_mask_hint()
+ ) from exc
+
+
def _selected_chat_template(tokenizer, tools: list | None) -> str:
"""Return the exact template that apply_chat_template will select.
@@ -196,18 +209,7 @@ def _selected_chat_template(tokenizer, tools: list | None) -> str:
its result rather than guessing from tokenizer.chat_template.
"""
label = _tokenizer_label(tokenizer)
- getter = getattr(tokenizer, "get_chat_template", None)
- try:
- if callable(getter):
- template = getter(chat_template=None, tools=tools)
- else:
- template = getattr(tokenizer, "chat_template", None)
- except Exception as exc:
- raise ExactAssistantMaskError(
- f"tokenizer {label!r} could not select its native chat template: {exc}. "
- + _legacy_mask_hint()
- ) from exc
-
+ template = _plain_selected_template(tokenizer, tools)
if not isinstance(template, str) or not template:
raise ExactAssistantMaskError(
f"tokenizer {label!r} has no native chat template, so an exact "
@@ -222,6 +224,59 @@ def _selected_chat_template(tokenizer, tools: list | None) -> str:
return template
+_QWEN_ASSISTANT_BRANCH = re.compile(
+ r'(\{%-? elif message\.role == "assistant" %\})(.*?)(\{%-? elif message\.role == "tool" %\})',
+ re.S,
+)
+_QWEN_TURN_OPEN = "'<|im_start|>' + message.role + '\\n"
+_QWEN_TURN_CLOSE = "{{- '<|im_end|>\\n' }}"
+_QWEN_THINK_DECISION = "{%- if loop.index0 > ns.last_query_index %}"
+
+
+def qwen_generation_template(template: str) -> str | None:
+ """Add `{% generation %}` to an official Qwen3 / Qwen3.5 chat template.
+
+ Qwen's published templates lack generation tags, so Transformers cannot
+ return an assistant mask for them. This rewrite touches only the assistant
+ branch: it emits `<|im_start|>assistant\\n` once before the existing
+ think/no-think decision, opens the generation block, strips that prefix
+ from every original emission, and closes the block after `<|im_end|>`.
+ The rendered text is therefore identical to the official template, and
+ the mask covers reasoning, content, tool calls and `<|im_end|>`. Any
+ template that does not match this exact Qwen shape returns None, so other
+ model families keep their old behavior (a clear ExactAssistantMaskError).
+ """
+ match = _QWEN_ASSISTANT_BRANCH.search(template)
+ if match is None:
+ return None
+ body = match.group(2)
+ if (
+ body.count(_QWEN_THINK_DECISION) != 1
+ or body.count(_QWEN_TURN_CLOSE) != 1
+ or "generation" in body
+ ):
+ return None
+ opens = body.count(_QWEN_TURN_OPEN)
+ if opens < 2:
+ return None
+ # "'<|im_start|>' + message.role + '\n' + content" -> "content";
+ # "'<|im_start|>' + message.role + '\n..." -> "'...".
+ body = body.replace(_QWEN_TURN_OPEN + "' + ", "")
+ body = body.replace(_QWEN_TURN_OPEN, "'")
+ if "<|im_start|>" in body:
+ return None
+ body = body.replace(
+ _QWEN_THINK_DECISION,
+ "{{- '<|im_start|>' + message.role + '\\n' }}\n {%- generation %}\n "
+ + _QWEN_THINK_DECISION,
+ )
+ body = body.replace(
+ _QWEN_TURN_CLOSE,
+ "{{- '<|im_end|>' }}\n {%- endgeneration %}\n {{- '\\n' }}",
+ )
+ return template[: match.start(2)] + body + template[match.end(2) :]
+
+
def _flat_list(value: Any, field: str, tokenizer) -> list:
if hasattr(value, "tolist"):
value = value.tolist()
@@ -258,12 +313,23 @@ def _native_assistant_tokens(
tokens are assistant targets. Altering either sequence after templating
would make the returned mask inexact.
"""
- _selected_chat_template(tokenizer, tools)
+ override = {}
+ try:
+ _selected_chat_template(tokenizer, tools)
+ except ExactAssistantMaskError as exc:
+ # Official Qwen3 / Qwen3.5 templates lack `{% generation %}`; use the
+ # render-identical tagged rewrite. Every other template still fails.
+ template = _plain_selected_template(tokenizer, tools)
+ patched = qwen_generation_template(template) if template else None
+ if patched is None:
+ raise
+ override = {"chat_template": patched}
label = _tokenizer_label(tokenizer)
try:
encoded = tokenizer.apply_chat_template(
messages,
tools=tools,
+ **override,
tokenize=True,
add_generation_prompt=False,
return_dict=True,