From 01841b283e795f3fab5b8ccababb0e61aa02ec4d Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Sun, 11 Oct 2026 16:07:27 +0000 Subject: [PATCH 1/4] codex_traces: keep Codex tool calls and tool outputs in session rows 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 (1 rollout: 21 msgs, 0 calls -> 161 msgs, 70 calls after this change). Co-Authored-By: Claude Opus 5.5 --- tests/test_codex_traces.py | 42 ++++++++++++++++++++++++++++++++++++++ yeto/codex_traces.py | 23 +++++++++++++++++++++ 2 files changed, 65 insertions(+) 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/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")) From fb3d510fbc9286d1f7dc35f003fa846c9d590fa6 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Sun, 11 Oct 2026 16:13:43 +0000 Subject: [PATCH 2/4] data: exact native assistant mask for official Qwen3/Qwen3.5 templates Official Qwen3 and Qwen3.5 chat templates lack {% generation %}, so --assistant-mask-mode native raised ExactAssistantMaskError. Rewrite only the assistant branch of a matching Qwen template into a render-identical tagged version and pass it as chat_template. The mask covers reasoning, content, tool calls and <|im_end|>; system, user and tool tokens stay 0. Templates that do not match the Qwen shape keep the existing error. Co-Authored-By: Claude Opus 5.5 --- tests/test_qwen_generation_template.py | 95 ++++++++++++++++++++++++++ yeto/data.py | 92 +++++++++++++++++++++---- 2 files changed, 174 insertions(+), 13 deletions(-) create mode 100644 tests/test_qwen_generation_template.py 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/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, From 9ea03e7ac802f21a9a996b6e44a78dc29de68309 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Sun, 11 Oct 2026 16:27:27 +0000 Subject: [PATCH 3/4] tools: add SFT held-out weighted-CE eval and adapter load self-check Co-Authored-By: Claude Opus 5.5 --- tools/eval_sft_heldout.py | 113 +++++++++++++++++++++++++++++++++ tools/test_eval_sft_heldout.py | 47 ++++++++++++++ 2 files changed, 160 insertions(+) create mode 100644 tools/eval_sft_heldout.py create mode 100644 tools/test_eval_sft_heldout.py 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"] From 8f952cca01996030e5adc1651c624b9685c047d5 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Sun, 11 Oct 2026 16:32:23 +0000 Subject: [PATCH 4/4] compare_diloco: give local arms a per-run island HMAC key The legacy syncer now refuses islands without YETO_ISLAND_HMAC_KEY, so every DiLoCo arm of the LM benchmark died at syncer start. All arms run on one host; a random per-run key in the harness environment is inherited by the syncer and every learner, and an operator-provided key is kept. Co-Authored-By: Claude Opus 5.5 --- scripts/compare_diloco.py | 11 +++++++++++ tests/test_lm_benchmark.py | 9 +++++++++ 2 files changed, 20 insertions(+) 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_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