From aed0d3f7886980deec745c10f99aa78b471b115a Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Fri, 13 Nov 2026 08:00:19 +0000 Subject: [PATCH 1/4] feat(aorlta): Qwen3 dense models besides Qwen3.5, from the tokenizer and the checkpoint's config The control-token ids and the generation prompts now come from the model's tokenizer and chat template (templates.ChatFormat), not Qwen3.5's constants: the rollout's stop token, the append rule's suffix check, the engines' self-test and the replies' stop stripping all use them. A parse adapter no longer has to name the stop token. The trainer loads the base with the transformers class of its family (model_family: Qwen3.5 or Qwen3, from config.json), the linear-attention fast-path check does not apply to a model with full attention only, and the LoRA target keys and the adapter merge tool accept Qwen3's model.layers.N names. Configs for the arithmetic example on Qwen3-0.6B. --- aorlta/configs/arith-score-qwen3-0.6b.yaml | 38 +++++++ aorlta/configs/arith-train-qwen3-0.6b.yaml | 43 ++++++++ aorlta/examples/hip-kernels/hip_kernels.py | 5 - aorlta/examples/qa/qa_task.py | 5 - aorlta/src/aorlta/model_family.py | 66 +++++++++++ aorlta/src/aorlta/orchestrator/wiring.py | 11 +- aorlta/src/aorlta/rollout/chain.py | 22 ++-- aorlta/src/aorlta/rollout/engine_pool.py | 10 +- aorlta/src/aorlta/rollout/feedback.py | 5 +- aorlta/src/aorlta/rollout/templates.py | 121 ++++++++++++++++++--- aorlta/src/aorlta/sync/merge.py | 10 +- aorlta/src/aorlta/sync/reload.py | 3 +- aorlta/src/aorlta/sync/writer.py | 19 +++- aorlta/src/aorlta/trainer/kernels.py | 38 ++++++- aorlta/src/aorlta/trainer/rank_pool.py | 9 +- aorlta/tests/trainer/test_kernels.py | 36 +++++- aorlta/tools/merge_base_adapter.py | 17 +-- 17 files changed, 391 insertions(+), 67 deletions(-) create mode 100644 aorlta/configs/arith-score-qwen3-0.6b.yaml create mode 100644 aorlta/configs/arith-train-qwen3-0.6b.yaml create mode 100644 aorlta/src/aorlta/model_family.py diff --git a/aorlta/configs/arith-score-qwen3-0.6b.yaml b/aorlta/configs/arith-score-qwen3-0.6b.yaml new file mode 100644 index 000000000..7b718d0f7 --- /dev/null +++ b/aorlta/configs/arith-score-qwen3-0.6b.yaml @@ -0,0 +1,38 @@ +# arith-score.yaml on Qwen/Qwen3-0.6B (a dense Qwen3, tied embeddings): the 64 held-out multiplications, 4 samples each. +# One GPU: one TokenSpeed engine on GPU 0 and, for training, one LoRA rank beside it; the answer check runs on the CPU. +# Site-specific paths, to edit here or to set in a site file (AORLTA_SITE): run.root, model.base_dir (and +# model.adapter_dir to score a trained adapter) and data.prompt_pools[0].path, which is this checkout's +# aorlta/examples/qa/arith/heldout. +plugins: [examples/qa/qa_task.py] +run: + run_id: arith-score-qwen3-0.6b + source: standin + root: /scratch/aorlta/runs +model: + base_dir: /data/hf/hub/models--Qwen--Qwen3-0.6B/snapshots/c1899de289a04d12100db370d81485cdf75e47ca + source: standin +data: + prompt_pools: + - path: /opt/aorlta/examples/qa/arith/heldout + source: standin + task: qa +rollout: + engines: [0] + K: 4 + P: 16 + max_repair_turns: 1 + thinking: "off" + temperature: 0.7 + max_new_tokens: 512 + compile_adapter: none + parse_adapter: qa.parse + feedback_template: qa.feedback + # a question with one short answer can get K identical replies: that is not a broken sampler here (the engines' + # start-up self-test checks that), so only stop when most groups are + identical_groups_stop: 0.75 +reward: + profile: qa-match + reward_fn: qa.v1 +orchestrator: + mode: score + run_seed: 20261113 diff --git a/aorlta/configs/arith-train-qwen3-0.6b.yaml b/aorlta/configs/arith-train-qwen3-0.6b.yaml new file mode 100644 index 000000000..d4271de34 --- /dev/null +++ b/aorlta/configs/arith-train-qwen3-0.6b.yaml @@ -0,0 +1,43 @@ +# arith-train.yaml on Qwen/Qwen3-0.6B (a dense Qwen3, tied embeddings): 512 generated 3-digit multiplications, 16 per +# iteration x 8 samples, 20 iterations. +# One GPU: one TokenSpeed engine on GPU 0 and, for training, one LoRA rank beside it; the answer check runs on the CPU. +# Site-specific paths, to edit here or to set in a site file (AORLTA_SITE): run.root, model.base_dir (and +# model.adapter_dir to score a trained adapter) and data.prompt_pools[0].path, which is this checkout's +# aorlta/examples/qa/arith/train. +plugins: [examples/qa/qa_task.py] +run: + run_id: arith-train-qwen3-0.6b + source: standin + root: /scratch/aorlta/runs +model: + base_dir: /data/hf/hub/models--Qwen--Qwen3-0.6B/snapshots/c1899de289a04d12100db370d81485cdf75e47ca + source: standin +data: + prompt_pools: + - path: /opt/aorlta/examples/qa/arith/train + source: standin + task: qa +rollout: + engines: [0] + K: 8 + P: 16 + max_repair_turns: 1 + thinking: "off" + temperature: 1.0 + max_new_tokens: 512 + compile_adapter: none + parse_adapter: qa.parse + feedback_template: qa.feedback + # a question with one short answer can get K identical replies: that is not a broken sampler here (the engines' + # start-up self-test checks that), so only stop when most groups are + identical_groups_stop: 0.75 +reward: + profile: qa-match + reward_fn: qa.v1 +orchestrator: + iterations: 20 + run_seed: 20261113 +trainer: + ranks: 1 + micro_batch_tokens: 20480 + lr: 2.0e-5 diff --git a/aorlta/examples/hip-kernels/hip_kernels.py b/aorlta/examples/hip-kernels/hip_kernels.py index 3a9a1a2be..174e520b5 100644 --- a/aorlta/examples/hip-kernels/hip_kernels.py +++ b/aorlta/examples/hip-kernels/hip_kernels.py @@ -96,11 +96,6 @@ class HipkParse: name = "hipk.parse" - def __init__(self): - from aorlta.rollout.templates import IM_END_ID - - self.reply_stops = (IM_END_ID,) - def check_prompts(self, prompts, thinking, max_repair_turns=0): for p in prompts: if not p.prompt_id.startswith("hipk:"): diff --git a/aorlta/examples/qa/qa_task.py b/aorlta/examples/qa/qa_task.py index e7bc9f709..0494b5ca5 100644 --- a/aorlta/examples/qa/qa_task.py +++ b/aorlta/examples/qa/qa_task.py @@ -88,11 +88,6 @@ class QaParse: name = "qa.parse" - def __init__(self): - from aorlta.rollout.templates import IM_END_ID - - self.reply_stops = (IM_END_ID,) - def check_prompts(self, prompts, thinking, max_repair_turns=0): return None diff --git a/aorlta/src/aorlta/model_family.py b/aorlta/src/aorlta/model_family.py new file mode 100644 index 000000000..94e9710c5 --- /dev/null +++ b/aorlta/src/aorlta/model_family.py @@ -0,0 +1,66 @@ +"""The model families AORLTA can train, told apart by a checkpoint's `config.json`. + +What differs between them outside the tokenizer (whose control-token ids and generation prompts +`rollout.templates.ChatFormat` reads from the tokenizer itself): + +- `qwen3_5`: Qwen3.5 (`Qwen3_5ForConditionalGeneration`). The checkpoints are vision-language, so transformers + loads them with `AutoModelForImageTextToText`, and the text decoder's keys are `model.language_model.layers.N.*`. + Its layers mix full attention with gated-DeltaNet linear attention, whose fast path the trainer requires + (`trainer.kernels`). +- `qwen3`: dense Qwen3 (`Qwen3ForCausalLM`, Qwen3-0.6B to Qwen3-32B). Loaded with `AutoModelForCausalLM`; the + decoder's keys are `model.layers.N.*`; full attention only, so the linear-attention fast path does not apply. + +A config whose `model_type` is neither, or whose `architectures` is not exactly the family's, is refused: the LoRA +targets, the name map and the trainer's loader are only known for these two. +""" + +from __future__ import annotations + +import json +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path + + +class ModelFamilyError(ValueError): + """A checkpoint config that names no family AORLTA supports, or cannot be read.""" + + +@dataclass(frozen=True) +class ModelFamily: + model_type: str + architecture: str # the one entry of config.json's `architectures` + auto_class: str # the transformers auto class that loads the checkpoint + decoder_prefix: str # the checkpoint keys of decoder layer N start with decoder_prefix + "N." + linear_attention: bool # whether some layers are gated-DeltaNet linear attention + + +FAMILIES = { + "qwen3_5": ModelFamily("qwen3_5", "Qwen3_5ForConditionalGeneration", "AutoModelForImageTextToText", + "model.language_model.layers.", True), + "qwen3": ModelFamily("qwen3", "Qwen3ForCausalLM", "AutoModelForCausalLM", "model.layers.", False), +} + + +def family_of(config: Mapping) -> ModelFamily: + """The family of a parsed `config.json`.""" + if not isinstance(config, Mapping): + raise ModelFamilyError("the model config is not a JSON object") + model_type = config.get("model_type") + fam = FAMILIES.get(model_type) if isinstance(model_type, str) else None + if fam is None: + raise ModelFamilyError(f"model_type {model_type!r} is not a supported family ({', '.join(FAMILIES)})") + archs = config.get("architectures") + if archs != [fam.architecture]: + raise ModelFamilyError(f"model_type {model_type!r} with architectures {archs!r}: expected [{fam.architecture!r}]") + return fam + + +def read_family(model_dir: str | Path) -> ModelFamily: + """The family of the checkpoint in `model_dir`, from its `config.json`.""" + path = Path(model_dir) / "config.json" + try: + config = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError) as exc: + raise ModelFamilyError(f"{path}: cannot be read: {type(exc).__name__}: {str(exc)[:120]}") from None + return family_of(config) diff --git a/aorlta/src/aorlta/orchestrator/wiring.py b/aorlta/src/aorlta/orchestrator/wiring.py index 512539093..43ece1697 100644 --- a/aorlta/src/aorlta/orchestrator/wiring.py +++ b/aorlta/src/aorlta/orchestrator/wiring.py @@ -330,13 +330,15 @@ class RolloutAdapter: def __init__(self, cfg: RunConfig, tokenizer, handles: Sequence[Any], *, max_len: int, spans, compiler, on_shutdown: Sequence[Callable[[], None]] = (), clock=None, selftest_messages: Sequence | None = None): from ..rollout.engine_pool import EnginePool + from ..rollout.templates import chat_format self.cfg, self.tok, self.spans, self.max_len, self.compiler = cfg, tokenizer, spans, max_len, compiler self.selftest_messages = None if selftest_messages is None else list(selftest_messages) # else the prompts self.clock = clock # the orchestrator's clock: its deadlines are read on it self.handles = [CountedHandle(h) for h in handles] self.pool = EnginePool(self.handles, run_seed=cfg.orchestrator.run_seed, - request_timeout_s=cfg.rollout.request_timeout_s, clock=spans.now) + request_timeout_s=cfg.rollout.request_timeout_s, clock=spans.now, + stop_token_ids=chat_format(tokenizer).stop_token_ids) self.manager = None self._on_shutdown = list(on_shutdown) @@ -1263,6 +1265,13 @@ def build(cfg: RunConfig, *, config_path=None, commit: str | None = None, parts: task = None if pool_cfg.task == "standin" else plugins.task(pool_cfg.task) if task is None and (parts.pool is None or parts.probes is None): # the prompts and probes come from the dataset dataset = standin_dataset_dir(pool_cfg.path) + if cfg.orchestrator.mode != "score" and parts.load_model is None: # the trainer loads the base by its family + from ..model_family import ModelFamilyError, read_family + + try: + read_family(cfg.model.base_dir) + except ModelFamilyError as exc: + raise WiringError(f"model.base_dir: {exc}") from None tok = parts.tokenizer or HFTokenizer.load(cfg.model.base_dir) if parts.pool is not None: pool = parts.pool diff --git a/aorlta/src/aorlta/rollout/chain.py b/aorlta/src/aorlta/rollout/chain.py index 4a3a28e83..fe3d3b312 100644 --- a/aorlta/src/aorlta/rollout/chain.py +++ b/aorlta/src/aorlta/rollout/chain.py @@ -60,7 +60,7 @@ from .engine_pool import EnginePool, RolloutStop, logprobs_aligned, lpt_assign from .feedback import Feedback, build_feedback from .templates import ( - IM_END_ID, STOP_TOKEN_IDS, TemplateError, Tokenizer, next_turn_ids, rerendered_turn_ids, turn0_ids, + IM_END_ID, STOP_TOKEN_IDS, TemplateError, Tokenizer, chat_format, next_turn_ids, rerendered_turn_ids, turn0_ids, ) SCORED_ENDS = frozenset({"compiled", "no_change", "no_code_block", "turns_exhausted", "context_limit", "logprob_error"}) @@ -78,7 +78,8 @@ def check(self, feedback: str, prompt: PromptRecord) -> bool: # True: the feedb class ParseAdapter(Protocol): name: str - reply_stops: tuple[int, ...] # the stop tokens dropped from the end of a reply before it is read + # optional: reply_stops, the stop tokens dropped from the end of a reply before it is read. Without it the + # model's own stop tokens are dropped; with it, it must hold them. def check_prompts(self, prompts: Sequence[PromptRecord], thinking: str, max_repair_turns: int) -> None: ... # raises ValueError @@ -101,7 +102,6 @@ class StandinParse: the chain (no repair turn).""" name = "standin.parse" - reply_stops = (IM_END_ID,) def check_prompts(self, prompts: Sequence[PromptRecord], thinking: str, max_repair_turns: int = 0) -> None: return None @@ -218,10 +218,11 @@ def from_config(cls, cfg: RunConfig, max_model_len: int) -> RolloutSettings: identical_groups_stop=r.identical_groups_stop, ) - def sampling_params(self, seed: int) -> dict: - """Section 6.1, Sampling: one request per chain per turn, with its own seed.""" + def sampling_params(self, seed: int, stop_token_ids: Sequence[int] = STOP_TOKEN_IDS) -> dict: + """Section 6.1, Sampling: one request per chain per turn, with its own seed, stopping on the model's stop + tokens.""" return {"temperature": self.temperature, "top_p": self.top_p, "top_k": self.top_k, - "max_new_tokens": self.max_new_tokens, "seed": seed, "stop_token_ids": list(STOP_TOKEN_IDS)} + "max_new_tokens": self.max_new_tokens, "seed": seed, "stop_token_ids": list(stop_token_ids)} @dataclass @@ -305,7 +306,12 @@ def __init__( raise RolloutStop(f"the {kind} given is {got!r}, the settings name {want!r}") if settings.repair_turn_ids not in REPAIR_TURN_IDS: raise RolloutStop(f"repair_turn_ids must be one of {REPAIR_TURN_IDS}, got {settings.repair_turn_ids!r}") - self.reply_stops = tuple(self.parse_adapter.reply_stops) + self.stop_token_ids = chat_format(tokenizer).stop_token_ids + stops = getattr(self.parse_adapter, "reply_stops", None) + self.reply_stops = self.stop_token_ids if stops is None else tuple(stops) + if not set(self.stop_token_ids) <= set(self.reply_stops): + raise RolloutStop(f"the parse adapter's reply_stops {self.reply_stops} do not hold the model's stop tokens " + f"{self.stop_token_ids}") self.spans = spans if spans is not None else SpanRecorder() self.compile_workers = compile_workers self.mean_response_tokens = ( @@ -366,7 +372,7 @@ def _run_chain(self, *, prompt: PromptRecord, group: int, ids0: list[int], k: in end, cut = "engine_error", True # the generation deadline: no further turn starts break turn_seed = seeds.seed_for(prompt.prompt_id, k, t) - sp = s.sampling_params(turn_seed) + sp = s.sampling_params(turn_seed, self.stop_token_ids) span = self.spans.start("gen.turn0" if t == 0 else "gen.repair", iteration, chain_id=cid, engine_id=engine, turn=t, input_len=len(ids)) out = self.pool.generate(engine, ids, sp) diff --git a/aorlta/src/aorlta/rollout/engine_pool.py b/aorlta/src/aorlta/rollout/engine_pool.py index 2cdb530aa..3b754e3e7 100644 --- a/aorlta/src/aorlta/rollout/engine_pool.py +++ b/aorlta/src/aorlta/rollout/engine_pool.py @@ -26,6 +26,7 @@ from typing import Any from ..ids import seed +from .templates import STOP_TOKEN_IDS SELF_TEST_PROMPTS = 4 SELF_TEST_TOKENS = 128 @@ -148,7 +149,8 @@ def to_dict(self) -> dict: "requests": self.requests, "seconds": round(self.seconds, 3), "errors": self.errors} -def run_self_test(engine: Any, prompts: Sequence[Sequence[int]], run_seed: int, timeout_s: float) -> SelfTest: +def run_self_test(engine: Any, prompts: Sequence[Sequence[int]], run_seed: int, timeout_s: float, + stop_token_ids: Sequence[int] = STOP_TOKEN_IDS) -> SelfTest: if len(prompts) < SELF_TEST_PROMPTS: raise ValueError(f"the self-test needs {SELF_TEST_PROMPTS} prompts, got {len(prompts)}") t0 = time.monotonic() @@ -158,7 +160,7 @@ def run_self_test(engine: Any, prompts: Sequence[Sequence[int]], run_seed: int, outs = [] for s in self_test_seeds(run_seed, i): # one at a time sp = {"temperature": SELF_TEST_TEMPERATURE, "top_p": SELF_TEST_TOP_P, "top_k": -1, - "max_new_tokens": SELF_TEST_TOKENS, "seed": s, "stop_token_ids": [248046]} + "max_new_tokens": SELF_TEST_TOKENS, "seed": s, "stop_token_ids": list(stop_token_ids)} res = engine.generate(list(ids), sp, timeout_s=timeout_s) requests += 1 if "error" in res: @@ -246,6 +248,7 @@ def __init__( on_restart: Callable[[int], bool] | None = None, make_handle: Callable[[int], Any] | None = None, clock: Callable[[], float] = time.time, + stop_token_ids: Sequence[int] = STOP_TOKEN_IDS, # the model's (templates.chat_format); Qwen3.5's by default ): if not handles: raise ValueError("an engine pool needs at least one engine") @@ -257,6 +260,7 @@ def __init__( self.on_restart = on_restart self.make_handle = make_handle self.clock = clock + self.stop_token_ids = tuple(stop_token_ids) self.prefix = [PrefixIndex() for _ in self.handles] self.available = [False] * len(self.handles) self.self_tests: list[dict] = [] @@ -281,7 +285,7 @@ def _bring_up(self, i: int, selftest_prompts: Sequence[Sequence[int]]) -> None: if problems: handle.stop() raise RolloutStop(f"engine {i} refused: " + "; ".join(problems)) - test = run_self_test(handle, selftest_prompts, self.run_seed, self.self_test_timeout_s) + test = run_self_test(handle, selftest_prompts, self.run_seed, self.self_test_timeout_s, self.stop_token_ids) with self._lock: self.self_tests.append(test.to_dict()) if not test.passed: diff --git a/aorlta/src/aorlta/rollout/feedback.py b/aorlta/src/aorlta/rollout/feedback.py index 645a624b1..62bb38b24 100644 --- a/aorlta/src/aorlta/rollout/feedback.py +++ b/aorlta/src/aorlta/rollout/feedback.py @@ -31,7 +31,7 @@ from collections.abc import Sequence from dataclasses import dataclass -from .templates import Tokenizer, repair_suffix +from .templates import Tokenizer, chat_format, repair_suffix HEADER = "Your code did not compile (hipcc --offload-arch=gfx950 -c). Errors:" FOOTER = "Fix the errors and reply with the complete corrected source in a single ```cpp code block." @@ -190,5 +190,6 @@ def numbered_list(feedback_text: str) -> str: def feedback_token_bound(tokenizer: Tokenizer, max_chars: int, thinking: str) -> int: """Tokens of the repair suffix around the longest list, for sizing `max_model_len`: the fixed text plus at most one token per character of the list [E; the per-chain length check before each submit is the exact guard].""" - fixed = tokenizer.encode(repair_suffix(HEADER + "\n" + MORE.format(n=99999) + "\n" + FOOTER, thinking)) + fixed = tokenizer.encode(repair_suffix(HEADER + "\n" + MORE.format(n=99999) + "\n" + FOOTER, thinking, + chat_format(tokenizer))) return len(fixed) + max_chars diff --git a/aorlta/src/aorlta/rollout/templates.py b/aorlta/src/aorlta/rollout/templates.py index 5598e4c88..b3f28830c 100644 --- a/aorlta/src/aorlta/rollout/templates.py +++ b/aorlta/src/aorlta/rollout/templates.py @@ -7,13 +7,18 @@ A repair turn is built one of two ways (`rollout.repair_turn_ids`). `append`, the stand-ins' default, is a token-level append: - ids_t = ids_{t-1} + reply_{t-1} + ([248046] if reply_{t-1} does not end in 248046 else []) + tok(SUFFIX_t) + ids_t = ids_{t-1} + reply_{t-1} + ([IM_END] if reply_{t-1} does not end in IM_END else []) + tok(SUFFIX_t) SUFFIX_t = "\\n<|im_start|>user\\n" + FEEDBACK_t + "<|im_end|>\\n<|im_start|>assistant\\n" + THINK_PREFIX `rerender` renders the whole conversation again with the chat template, each reply as an assistant turn (its decoded text) and each feedback as a user turn -(`rerendered_turn_ids`). The two differ: Qwen3.5's template drops the think block of every assistant turn before -the last user turn, trims every turn, and re-tokenises the reply's text. C1.T2 pins the first difference. +(`rerendered_turn_ids`). The two differ: Qwen3.5's template (and Qwen3's) drops the think block of every assistant +turn before the last user turn, trims every turn, and re-tokenises the reply's text. C1.T2 pins the first difference. + +The control-token ids and THINK_PREFIX are the model's (`ChatFormat`), read from its tokenizer and chat template: +IM_END is 248046 for Qwen3.5 and 151645 for Qwen3, and with thinking on Qwen3.5's generation prompt opens the think +block (`\\n`) while Qwen3's leaves it to the model. A tokenizer that does not carry a format of its own (the +tests' stand-ins) is taken to have Qwen3.5's. The stand-in prompts are the original listings of a kernel dataset (`kernels.json` and `listings/.json`; `examples/standin-pool/` shows the format) plus a plain "fix this kernel" instruction (Section 6.10 item 7). @@ -25,12 +30,14 @@ import json import os from collections.abc import Mapping, Sequence +from dataclasses import dataclass from pathlib import Path from typing import Any, Protocol from ..ids import sha256_file, sha256_ids, standin_prompt_id from ..schemas import PromptRecord, TaskRef +# Qwen3.5's control-token ids (`QWEN3_5` below); a run uses its own model's (`chat_format`). ENDOFTEXT_ID = 248044 IM_START_ID = 248045 IM_END_ID = 248046 # <|im_end|>: the engine stops on it (Appendix C) @@ -41,6 +48,12 @@ STOP_TOKEN_IDS = (IM_END_ID,) THINK_PREFIX = {"off": "\n\n\n\n", "on": "\n"} +ASSISTANT_OPEN = "<|im_start|>assistant\n" +# The generation prompt after ASSISTANT_OPEN that a chat template may produce for each thinking flag: thinking off +# must close an empty think block, or the flag would not be honoured; thinking on may open the block or leave it to +# the model. +THINK_PREFIXES_KNOWN = {"off": ("\n\n\n\n",), "on": ("\n", "")} +_PROBE_CHAT = ({"role": "system", "content": "s"}, {"role": "user", "content": "u"}) # The stand-in texts (the earlier prototype's, unchanged). STANDIN_SYSTEM = "You are an expert AMD GPU kernel engineer." @@ -58,6 +71,69 @@ class TemplateError(ValueError): """A turn sequence cannot be built exactly as Appendix C specifies.""" +@dataclass(frozen=True) +class ChatFormat: + """A model's five control tokens and the think prefix its chat template puts after ASSISTANT_OPEN for each + thinking flag.""" + + endoftext_id: int + im_start_id: int + im_end_id: int + think_open_id: int + think_close_id: int + think_prefix_off: str + think_prefix_on: str + + @property + def control_ids(self) -> frozenset[int]: + return frozenset({self.endoftext_id, self.im_start_id, self.im_end_id, self.think_open_id, + self.think_close_id}) + + @property + def stop_token_ids(self) -> tuple[int, ...]: + return (self.im_end_id,) + + def think_prefix(self, thinking: str) -> str: + _thinking_on(thinking) + return self.think_prefix_off if thinking == "off" else self.think_prefix_on + + @classmethod + def derive(cls, tokenizer: Tokenizer) -> ChatFormat: + """Read the format from a tokenizer: each control text must encode to exactly one id, the five ids distinct, + and the chat template's generation prompt for each thinking flag must be ASSISTANT_OPEN followed by one of + THINK_PREFIXES_KNOWN. Anything else is refused: the turn rules below are only right for such a model.""" + ids = {} + for text in CONTROL_TEXTS: + enc = tokenizer.encode(text) + if len(enc) != 1: + raise TemplateError(f"the tokenizer encodes {text!r} to {len(enc)} ids, not one control token") + ids[text] = int(enc[0]) + if len(set(ids.values())) != len(ids): + raise TemplateError(f"the tokenizer's control tokens are not distinct: {ids}") + prefixes = {} + for thinking in ("off", "on"): + base = tokenizer.render(list(_PROBE_CHAT), thinking, add_generation_prompt=False) + full = tokenizer.render(list(_PROBE_CHAT), thinking, add_generation_prompt=True) + gen = full[len(base):] if full.startswith(base) else None + prefix = gen[len(ASSISTANT_OPEN):] if gen is not None and gen.startswith(ASSISTANT_OPEN) else None + if prefix not in THINK_PREFIXES_KNOWN[thinking]: + raise TemplateError(f"the chat template's thinking-{thinking} generation prompt is {gen!r}, not " + f"{ASSISTANT_OPEN!r} + one of {THINK_PREFIXES_KNOWN[thinking]}") + prefixes[thinking] = prefix + return cls(ids["<|endoftext|>"], ids["<|im_start|>"], ids["<|im_end|>"], ids[""], ids[""], + prefixes["off"], prefixes["on"]) + + +QWEN3_5 = ChatFormat(ENDOFTEXT_ID, IM_START_ID, IM_END_ID, THINK_OPEN_ID, THINK_CLOSE_ID, THINK_PREFIX["off"], + THINK_PREFIX["on"]) + + +def chat_format(tokenizer: Tokenizer) -> ChatFormat: + """The tokenizer's own format (`HFTokenizer` reads it from the model's tokenizer), else Qwen3.5's.""" + fmt = getattr(tokenizer, "chat_format", None) + return QWEN3_5 if fmt is None else fmt + + class Tokenizer(Protocol): """What the rollout needs from the model's tokenizer.""" @@ -76,6 +152,14 @@ class HFTokenizer: def __init__(self, tok: Any): self.tok = tok + self._format: ChatFormat | None = None + + @property + def chat_format(self) -> ChatFormat: + """The model's control tokens and generation prompts, read from this tokenizer (`ChatFormat.derive`).""" + if self._format is None: + self._format = ChatFormat.derive(self) + return self._format @classmethod def load(cls, model_dir: str | Path) -> HFTokenizer: @@ -108,9 +192,8 @@ def _thinking_on(thinking: str) -> bool: return thinking == "on" -def generation_prompt(thinking: str) -> str: - _thinking_on(thinking) - return "<|im_start|>assistant\n" + THINK_PREFIX[thinking] +def generation_prompt(thinking: str, fmt: ChatFormat = QWEN3_5) -> str: + return ASSISTANT_OPEN + fmt.think_prefix(thinking) def _check_messages(messages: Any) -> list[dict]: @@ -140,7 +223,7 @@ def render_turn0(tokenizer: Tokenizer, messages: Sequence[Mapping[str, Any]], th """`apply_chat_template(messages, add_generation_prompt=True, enable_thinking=...)`, tokenised with add_special_tokens=False. The ids must end in the generation prompt for the thinking flag.""" msgs = _check_messages(messages) - tail = tokenizer.encode(generation_prompt(thinking)) + tail = tokenizer.encode(generation_prompt(thinking, chat_format(tokenizer))) ids = tokenizer.encode(tokenizer.render(msgs, thinking)) if len(ids) <= len(tail) or ids[-len(tail):] != tail: raise TemplateError(f"the rendered turn 0 does not end in the thinking-{thinking} generation prompt") @@ -164,13 +247,14 @@ def turn0_ids(prompt: PromptRecord, tokenizer: Tokenizer, thinking: str) -> list return ids -def repair_suffix(feedback: str, thinking: str) -> str: - return "\n<|im_start|>user\n" + feedback + "<|im_end|>\n" + generation_prompt(thinking) +def repair_suffix(feedback: str, thinking: str, fmt: ChatFormat = QWEN3_5) -> str: + return "\n<|im_start|>user\n" + feedback + "<|im_end|>\n" + generation_prompt(thinking, fmt) -def _expected_suffix_controls(thinking: str) -> list[int]: - think = [THINK_OPEN_ID, THINK_CLOSE_ID] if thinking == "off" else [THINK_OPEN_ID] - return [IM_START_ID, IM_END_ID, IM_START_ID, *think] +def _expected_suffix_controls(thinking: str, fmt: ChatFormat) -> list[int]: + known = {"": fmt.think_open_id, "": fmt.think_close_id} + think = [known[t] for t in ("", "") if t in fmt.think_prefix(thinking)] + return [fmt.im_start_id, fmt.im_end_id, fmt.im_start_id, *think] def next_turn_ids( @@ -183,12 +267,13 @@ def next_turn_ids( raise TemplateError("feedback must be a non-empty str") if any(t in feedback for t in CONTROL_TEXTS): raise TemplateError("feedback contains the text of a control token") + fmt = chat_format(tokenizer) ids = [int(t) for t in prev_ids] + [int(t) for t in reply_ids] - if not reply_ids or int(reply_ids[-1]) != IM_END_ID: - ids.append(IM_END_ID) - suffix = tokenizer.encode(repair_suffix(feedback, thinking)) - controls = [t for t in suffix if t in CONTROL_IDS] - if controls != _expected_suffix_controls(thinking): + if not reply_ids or int(reply_ids[-1]) != fmt.im_end_id: + ids.append(fmt.im_end_id) + suffix = tokenizer.encode(repair_suffix(feedback, thinking, fmt)) + controls = [t for t in suffix if t in fmt.control_ids] + if controls != _expected_suffix_controls(thinking, fmt): raise TemplateError(f"the repair suffix tokenised to control ids {controls}") return ids + suffix @@ -214,7 +299,7 @@ def rerendered_turn_ids( if not isinstance(reply, str): raise TemplateError("the reply must be a str") msgs = _check_messages([*messages, {"role": "assistant", "content": reply}, {"role": "user", "content": feedback}]) - tail = tokenizer.encode(generation_prompt(thinking)) + tail = tokenizer.encode(generation_prompt(thinking, chat_format(tokenizer))) ids = tokenizer.encode(tokenizer.render(msgs, thinking)) if len(ids) <= len(tail) or ids[-len(tail):] != tail: raise TemplateError(f"the rebuilt repair turn does not end in the thinking-{thinking} generation prompt") diff --git a/aorlta/src/aorlta/sync/merge.py b/aorlta/src/aorlta/sync/merge.py index 8e92392db..48f6ff69e 100644 --- a/aorlta/src/aorlta/sync/merge.py +++ b/aorlta/src/aorlta/sync/merge.py @@ -36,14 +36,18 @@ ROUNDINGS = ("nearest", "stochastic") ROUNDING_SEED_TAG = "aorlta.stochastic_rounding/1" -# PEFT target modules (the prototype's TARGETS) and the checkpoint keys they map to (Section 6.4, Model). +# PEFT target modules (the prototype's TARGETS) and the checkpoint keys they map to (Section 6.4, Model): the +# attention and MLP projections of every text decoder layer, and in Qwen3.5 the linear-attention projections too (a +# Qwen3 layer has none, so those names match nothing there). The decoder's layers are `model.language_model.layers.N` +# in a Qwen3.5 checkpoint and `model.layers.N` in a Qwen3 one (`model_family`); Qwen3.5's vision tower and its MTP +# head (`mtp.layers.0`) are never targets. TARGET_MODULE_PATTERN = ( r".*layers\.\d+\.(self_attn\.(q_proj|k_proj|v_proj|o_proj)|mlp\.(gate_proj|up_proj|down_proj)" r"|linear_attn\.(in_proj_qkv|in_proj_z|in_proj_a|in_proj_b|out_proj))" ) TARGET_KEY_RE = re.compile( - r".*language_model\.layers\.\d+\.(self_attn\.(q_proj|k_proj|v_proj|o_proj)|mlp\.(gate_proj|up_proj|down_proj)" - r"|linear_attn\.(in_proj_qkv|in_proj_z|in_proj_a|in_proj_b|out_proj))\.weight" + r"model\.(?:language_model\.)?layers\.\d+\.(self_attn\.(q_proj|k_proj|v_proj|o_proj)" + r"|mlp\.(gate_proj|up_proj|down_proj)|linear_attn\.(in_proj_qkv|in_proj_z|in_proj_a|in_proj_b|out_proj))\.weight" ) GROUP_RE = re.compile(r"\.(self_attn|mlp|linear_attn)\.([a-z_]+)\.weight$") diff --git a/aorlta/src/aorlta/sync/reload.py b/aorlta/src/aorlta/sync/reload.py index 91ee77ea3..7b2608718 100644 --- a/aorlta/src/aorlta/sync/reload.py +++ b/aorlta/src/aorlta/sync/reload.py @@ -2,7 +2,8 @@ `update_weights_from_disk(path, flush_cache=True)` replies `(success, message, num_paused_requests)`, and on success the message ends "(N checkpoint tensors read)", N counting the tensors the loader yielded (A5). A reply -counts only if success is exactly True and N is exactly the slot's tensor count (248 partial, 738 full). +counts only if success is exactly True and N is exactly the slot's tensor count (`writer.expected_count`: for +Qwen3.5-4B 248 partial, 738 full). A flushing reload is refused while a request pins the prefix cache (A7), so reloads happen only behind the generation barrier; a reload attempted with requests in flight is a programming error and stops the run. diff --git a/aorlta/src/aorlta/sync/writer.py b/aorlta/src/aorlta/sync/writer.py index 352e50efa..bfc248473 100644 --- a/aorlta/src/aorlta/sync/writer.py +++ b/aorlta/src/aorlta/sync/writer.py @@ -1,9 +1,11 @@ """C5 steps 2, 3 and 7: the name map, the slot writer, slot cleanup (design Section 6.5). TokenSpeed's Qwen3.5 `load_weights` silently skips a tensor name it does not know, yet counts it as read -(Appendix A, A5/A6). So every name written must be a base-index key, and the written set must be exactly the -target keys (partial checkpoint, D11) or exactly every base key (full checkpoint). Anything else stops the run -before a byte is written. +(Appendix A, A5/A6). Its Qwen3 `load_weights` fails on an unknown name instead, and with tied embeddings +(Qwen3-0.6B to 4B) skips the checkpoint's `lm_head.weight` but counts it too: a full checkpoint writes that tensor +unchanged, since no LoRA target touches it. So every name written must be a base-index key, and the written set +must be exactly the target keys (partial checkpoint, D11) or exactly every base key (full checkpoint). Anything +else stops the run before a byte is written. A slot is `//w{n mod slots}`: safetensors shards of at most 4 GiB, a `model.safetensors.index.json` for those names, the base directory's config and tokenizer files, and @@ -93,7 +95,13 @@ def read_base_index(base_dir: str | Path) -> dict[str, str]: def target_keys(base_keys: Iterable[str]) -> frozenset[str]: - return frozenset(k for k in base_keys if TARGET_KEY_RE.fullmatch(k)) + """The base keys that are LoRA targets (`merge.TARGET_KEY_RE`), all under one decoder prefix: a base holding + both `model.language_model.layers.N` and `model.layers.N` targets is refused.""" + targets = frozenset(k for k in base_keys if TARGET_KEY_RE.fullmatch(k)) + prefixes = sorted({_LAYER_RE.split(k, 1)[0] for k in targets}) + if len(prefixes) > 1: + raise NameMapError(f"the base's target keys are under more than one decoder prefix: {prefixes}") + return targets def module_key(module_name: str, targets: frozenset[str]) -> str: @@ -136,7 +144,8 @@ def name_map(module_names: Iterable[str], base_keys: Iterable[str]) -> dict[str, def expected_count(base_keys: Iterable[str], *, partial: bool) -> int: - """The tensor count a reload of a slot must report: 248 (partial) or 738 (full) for Qwen3.5-4B.""" + """The tensor count a reload of a slot must report, from the base's own keys: 248 (partial) or 738 (full) for + Qwen3.5-4B, 196 or 311 for Qwen3-0.6B, 252 or 399 for Qwen3-8B.""" base_keys = list(base_keys) return len(target_keys(base_keys)) if partial else len(set(base_keys)) diff --git a/aorlta/src/aorlta/trainer/kernels.py b/aorlta/src/aorlta/trainer/kernels.py index f87f238b0..6aacb3143 100644 --- a/aorlta/src/aorlta/trainer/kernels.py +++ b/aorlta/src/aorlta/trainer/kernels.py @@ -8,6 +8,11 @@ Detection is fail-closed: an import error, an unknown transformers layout, a model with no (or too few) linear-attention layers, or a probe that raises all count as "not available". + +The check applies only to a model that has linear attention. A model whose config declares its layer types and +none of them `linear_attention`, and which holds no linear-attention module (dense Qwen3, say), runs full attention +only: there is no fast path to miss, and the report says the check does not apply (`applies` False). A model whose +config does not declare its layer types is checked as before. """ from __future__ import annotations @@ -61,15 +66,19 @@ class FastPathReport: components: dict = field(default_factory=dict) linear_attention_layers: int | None = None versions: dict = field(default_factory=dict) + applies: bool = True # False: the model has no linear attention, so there is no fast path to check def to_dict(self) -> dict: - return { + out = { "available": self.available, "reasons": list(self.reasons), "components": dict(self.components), "linear_attention_layers": self.linear_attention_layers, "versions": dict(self.versions), } + if not self.applies: + out["applies"] = False + return out class FastPathUnavailable(RuntimeError): @@ -119,6 +128,18 @@ def _expected_linear_layers(model) -> int | None: return None +def without_linear_attention(model) -> bool: + """True only when the model positively has no linear attention: its config (or text config) lists its layer + types and none is `linear_attention`, and no module is a linear-attention layer (named `linear_attn`, or of a + class whose name holds `GatedDeltaNet`). Anything less certain is False, so the fast-path check runs.""" + if _expected_linear_layers(model) != 0: + return False + for name, m in model.named_modules(): + if name.rsplit(".", 1)[-1] == "linear_attn" or "GatedDeltaNet" in type(m).__name__: + return False + return True + + def _torch_path(layer, mod) -> list[str]: """The kernels `layer` runs on transformers' torch path (a missing one included).""" bad = [] @@ -154,6 +175,13 @@ def _model_reasons(model, mod) -> tuple[list[str], int]: def _check(model, modeling) -> FastPathReport: versions = installed_versions() + if model is not None: + try: + dense = without_linear_attention(model) + except Exception as exc: + return FastPathReport(False, (f"fast-path probe failed: {exc!r}",), versions=versions) + if dense: + return FastPathReport(True, (), {}, 0, versions, applies=False) try: mod = modeling if modeling is not None else importlib.import_module(MODELING) except Exception as exc: @@ -182,8 +210,11 @@ def check_model(model, modeling=None) -> FastPathReport: def honours_packing(model, modeling=None) -> bool: """Whether a packed row's sequences stay apart in `model`'s linear-attention layers: no layer runs one of `PACKING_KERNELS` on the torch path, which runs the row as one sequence (each sequence after the first starts - from the state the earlier ones left). Fail-closed: an import error or an unknown layout counts as not.""" + from the state the earlier ones left). A model without linear attention (`without_linear_attention`) does. + Fail-closed: an import error or an unknown layout counts as not.""" try: + if without_linear_attention(model): + return True mod = modeling if modeling is not None else importlib.import_module(MODELING) if not (hasattr(mod, "Qwen3_5GatedDeltaNet") and hasattr(mod, TORCH_FALLBACKS["chunk_gated_delta_rule"])): return False @@ -237,6 +268,9 @@ def ensure_fast_path( ) -> FastPathReport: """Check the fast path; raise `FastPathUnavailable` if it is missing and `require_fast_path` is set.""" report = check_model(model, modeling) if model is not None else check_environment(modeling) + if not report.applies: + log.info("the model has no linear-attention layers: the fast-path check does not apply") + return report if watcher is not None and watcher.fallback_logged: report = FastPathReport( False, diff --git a/aorlta/src/aorlta/trainer/rank_pool.py b/aorlta/src/aorlta/trainer/rank_pool.py index 1b3ed148f..c3a111d78 100644 --- a/aorlta/src/aorlta/trainer/rank_pool.py +++ b/aorlta/src/aorlta/trainer/rank_pool.py @@ -123,13 +123,16 @@ def file_init_method(directory) -> str: def load_policy_model(base_dir: str, device: str): - """The bf16 base as WP4 loads it, with the log watcher that `prepare_model` reads.""" - from transformers import AutoModelForImageTextToText + """The bf16 base as WP4 loads it, with the log watcher that `prepare_model` reads. The transformers class is the + checkpoint's family's (`model_family`): AutoModelForImageTextToText for Qwen3.5, AutoModelForCausalLM for Qwen3.""" + import transformers + from aorlta.model_family import read_family from aorlta.trainer.kernels import FallbackLogWatcher + auto = getattr(transformers, read_family(base_dir).auto_class) with FallbackLogWatcher() as watcher: - model = AutoModelForImageTextToText.from_pretrained(base_dir, dtype=torch.bfloat16, local_files_only=True) + model = auto.from_pretrained(base_dir, dtype=torch.bfloat16, local_files_only=True) model.to(device) return model, watcher diff --git a/aorlta/tests/trainer/test_kernels.py b/aorlta/tests/trainer/test_kernels.py index aaaed0ae3..a856dc0f6 100644 --- a/aorlta/tests/trainer/test_kernels.py +++ b/aorlta/tests/trainer/test_kernels.py @@ -203,12 +203,40 @@ def test_layer_with_a_missing_kernel_attribute_is_caught(): assert not kernels.check_model(model, mod).available -def test_model_without_linear_attention_layers_is_not_available(): +@pytest.mark.parametrize("layer_types", [["linear_attention"] * 3 + ["full_attention"], None]) +def test_model_without_linear_attention_layers_is_not_available(layer_types): + # its config declares linear-attention layers, or does not say: a model where none is found is refused mod = fake_modeling() - rep = kernels.check_model(FakeModel([], layer_types=["full_attention"] * 4), mod) - assert not rep.available + rep = kernels.check_model(FakeModel([], layer_types=layer_types), mod) + assert not rep.available and rep.applies assert rep.linear_attention_layers == 0 assert any("no Qwen3.5 linear-attention layers" in r for r in rep.reasons) + with pytest.raises(kernels.FastPathUnavailable): + kernels.ensure_fast_path(True, FakeModel([], layer_types=layer_types), modeling=mod) + + +@pytest.mark.parametrize("nested", [True, False]) +def test_the_check_does_not_apply_to_a_model_with_full_attention_only(nested): + # dense Qwen3: the config lists its layer types, none linear, and no module is linear attention + model = FakeModel([], layer_types=["full_attention"] * 4, nested=nested) + for mod in (fake_modeling(), fake_modeling(conv=False, fla=False)): + rep = kernels.ensure_fast_path(True, model, modeling=mod) + assert rep.available and not rep.applies and rep.reasons == () + assert rep.to_dict()["applies"] is False + assert kernels.without_linear_attention(model) + assert kernels.honours_packing(model, fake_modeling(conv=False, fla=False)) + watcher = kernels.FallbackLogWatcher() + watcher.fallback_logged = True # a Qwen3.5 build elsewhere in the process logged it: it says nothing here + assert kernels.ensure_fast_path(True, model, modeling=fake_modeling(), watcher=watcher).available + + +def test_a_linear_attention_module_keeps_the_check_whatever_the_config_says(): + mod = fake_modeling(fla=False) + model = FakeModel([mod.Qwen3_5GatedDeltaNet()], layer_types=["full_attention"] * 4) # config and modules disagree + assert not kernels.without_linear_attention(model) + with pytest.raises(kernels.FastPathUnavailable): + kernels.ensure_fast_path(True, model, modeling=mod) + assert kernels.check_model(fast_model(fake_modeling()), fake_modeling()).to_dict().get("applies", True) def test_layer_count_must_match_the_config(): @@ -367,7 +395,7 @@ def test_an_unknown_layout_does_not_honour_packing(name): def test_an_import_failure_does_not_honour_packing(monkeypatch): monkeypatch.setattr(kernels, "MODELING", "aorlta_no_such_package.modeling_qwen3_5") - assert not kernels.honours_packing(FakeModel([])) + assert not kernels.honours_packing(FakeModel([object()])) # a linear_attn module: the check applies # --- the real module --------------------------------------------------------------------------------------- diff --git a/aorlta/tools/merge_base_adapter.py b/aorlta/tools/merge_base_adapter.py index 921a2fcd7..0343823ab 100644 --- a/aorlta/tools/merge_base_adapter.py +++ b/aorlta/tools/merge_base_adapter.py @@ -4,10 +4,11 @@ For every adapter module, W' = bf16(W.float() + (alpha/r) * B.float() @ A.float()), in fp32 on the CPU, cast with round-to-nearest-even. That is `aorlta.sync.merge.merge_all`, the formula the trainer and the weight sync use. The modules go through WP3's name map (`aorlta.sync.writer.name_map`), so the merged set is exactly the base's LoRA target -keys (248 for Qwen3.5-4B). On top of it, an adapter module must be named `model.layers.N.` (the text-only -model) or `model.language_model.layers.N.` (the checkpoint), and nothing but lora_A/lora_B pairs may be in the -adapter: an embedding, head or other tensor is refused, as is any adapter_config field the plain formula does not -cover (DoRA, rsLoRA, rank or alpha patterns, modules_to_save, ...). +keys (248 for Qwen3.5-4B, 196 for Qwen3-0.6B). On top of it, an adapter module must be named `model.layers.N.` +or `model.language_model.layers.N.`, and map to the base's own key for that layer (under +`model.language_model.` in a Qwen3.5 checkpoint, `model.` in a Qwen3 one). Nothing but lora_A/lora_B pairs may be in +the adapter: an embedding, head or other tensor is refused, as is any adapter_config field the plain formula does +not cover (DoRA, rsLoRA, rank or alpha patterns, modules_to_save, ...). The output keeps the base's layout byte for byte. Every base file is copied, and only the target tensors' data ranges inside the shards are rewritten, so the shard names, the safetensors headers, the tensor offsets and the index are @@ -224,15 +225,17 @@ def load_adapter(path: str | Path) -> Adapter: def checkpoint_keys(adapter: Adapter, base_keys: list[str]) -> dict[str, str]: """module -> base checkpoint key, through WP3's name map (exactly the base's target keys, one to one), and the - stricter rule that the key is `model.language_model.` + the module's `layers.N.` + `.weight`.""" + stricter rule that the key is the base's decoder prefix (`model.language_model.` in a Qwen3.5 checkpoint, + `model.` in a Qwen3 one) + the module's `layers.N.` + `.weight`.""" try: mapping = name_map(adapter.modules, base_keys) except SyncError as exc: raise Refused(f"name map: {exc}") from exc + prefix = "model.language_model." if any(k.startswith("model.language_model.") for k in mapping.values()) else "model." for module, key in mapping.items(): tail = MODULE_RE.fullmatch(module).group("tail") - if key != f"model.language_model.{tail}.weight": - raise Refused(f"{module} mapped to {key}, not model.language_model.{tail}.weight") + if key != f"{prefix}{tail}.weight": + raise Refused(f"{module} mapped to {key}, not {prefix}{tail}.weight") return mapping From 9c9b0c5eac24a59f5e7a2cd26ffd00eb28ca1461 Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Fri, 13 Nov 2026 10:22:54 +0000 Subject: [PATCH 2/4] test(aorlta): the wired loop end to end on CPU on a tiny dense Qwen3 A tiny Qwen3 with Qwen3's control ids, its chat template and Qwen3-0.6B's checkpoint layout (one safetensors file, tied embeddings with lm_head stored) goes through the rollout, the real trainer loader, the name map, the merge, the slot writer and the reload. The fake engine stops on the request's stop token, as TokenSpeed does, rather than Qwen3.5's. --- aorlta/src/aorlta/rollout/fake_engine.py | 12 +- aorlta/tests/orchestrator/test_wiring_e2e.py | 4 +- .../orchestrator/test_wiring_e2e_qwen3.py | 188 ++++++++++++++++++ 3 files changed, 197 insertions(+), 7 deletions(-) create mode 100644 aorlta/tests/orchestrator/test_wiring_e2e_qwen3.py diff --git a/aorlta/src/aorlta/rollout/fake_engine.py b/aorlta/src/aorlta/rollout/fake_engine.py index d89a0b4e6..7168d95e2 100644 --- a/aorlta/src/aorlta/rollout/fake_engine.py +++ b/aorlta/src/aorlta/rollout/fake_engine.py @@ -3,7 +3,8 @@ `FakeEngine.generate` takes the arguments AORLTA passes to TokenSpeed's `Engine.generate` and returns the same shape: `{"text", "output_ids", "meta_info"}`, with `meta_info` carrying `finish_reason` ({"type": "stop", "matched": 248046} or {"type": "length", "length": N}, as -measured on the engine), token counts, `cached_tokens` and, when asked for and enabled, +measured on the engine; the stop is the request's first `stop_token_ids` entry, 248046 if it names none), token +counts, `cached_tokens` and, when asked for and enabled, `output_token_logprobs` as (logprob, token_id, None) triples, which is what TokenSpeed returns at 3b684074 (`engine/logprobs.py:255-259`). @@ -307,9 +308,10 @@ def log(outcome: str, completion_tokens: int = 0) -> None: content = self.tokenizer.encode(self._answer_text(rng)) max_new = sp["max_new_tokens"] + stop_id = (sp.get("stop_token_ids") or [STOP_TOKEN_ID])[0] # the model's stop token, as the request names it if len(content) + 1 <= max_new: - output_ids = content + [STOP_TOKEN_ID] - finish = {"type": "stop", "matched": STOP_TOKEN_ID} + output_ids = content + [stop_id] + finish = {"type": "stop", "matched": stop_id} else: output_ids = content[:max_new] finish = {"type": "length", "length": max_new} @@ -328,8 +330,8 @@ def log(outcome: str, completion_tokens: int = 0) -> None: meta["output_token_logprobs"] = self._logprobs(rng, output_ids) log("ok", len(output_ids)) self._wait(latency) - return {"text": self.tokenizer.decode(output_ids, skip_special_tokens=True), "output_ids": output_ids, - "meta_info": meta} + text = self.tokenizer.decode(output_ids[:-1] if finish["type"] == "stop" else output_ids, skip_special_tokens=True) + return {"text": text, "output_ids": output_ids, "meta_info": meta} def _check_sampling(self, sampling_params: Any) -> dict: if not isinstance(sampling_params, dict): diff --git a/aorlta/tests/orchestrator/test_wiring_e2e.py b/aorlta/tests/orchestrator/test_wiring_e2e.py index 1983c23b5..2c1cf63c3 100644 --- a/aorlta/tests/orchestrator/test_wiring_e2e.py +++ b/aorlta/tests/orchestrator/test_wiring_e2e.py @@ -99,10 +99,10 @@ class TwinEngine: CHUNK = 32 - def __init__(self, base_dir: Path, fake: FakeEngineConfig, engine_id: int): + def __init__(self, base_dir: Path, fake: FakeEngineConfig, engine_id: int, loader=None): self.fake = FakeEngine(fake, engine_id=engine_id, **ENGINE_KW) self.tokenizer_manager = self.fake.tokenizer_manager - self.model = load_tiny(base_dir) + self.model = (loader or load_tiny)(base_dir) self.params = dict(self.model.named_parameters()) self.reloads: list[str] = [] self._lock = threading.Lock() diff --git a/aorlta/tests/orchestrator/test_wiring_e2e_qwen3.py b/aorlta/tests/orchestrator/test_wiring_e2e_qwen3.py new file mode 100644 index 000000000..71140fb29 --- /dev/null +++ b/aorlta/tests/orchestrator/test_wiring_e2e_qwen3.py @@ -0,0 +1,188 @@ +"""The wired components end to end on CPU on a tiny dense Qwen3, as test_wiring_e2e runs them on a tiny Qwen3.5. + +What a Qwen3 changes, and this run goes through: + - the tokenizer's control ids (Qwen3's 151643-151668) and its chat template, whose thinking-on generation prompt + opens no think block: the rollout's stop token, the repair suffix and the replies' stop stripping come from it; + - the trainer's loader: the real one (`wiring.load_policy_model`), which picks AutoModelForCausalLM from + config.json, with `trainer.require_fast_path` on and the linear-attention check not applying; + - the decoder's `model.layers.N` names, through the name map, the merge, the slot writer and the reload; + - tied embeddings stored as Qwen3-0.6B stores them: one `model.safetensors` with no index, holding + `lm_head.weight` beside `model.embed_tokens.weight`. A full slot carries it unchanged; the engine replica skips + it (TokenSpeed's Qwen3 does, with tied embeddings) and counts it. +""" + +from __future__ import annotations + +import re +import sys +from pathlib import Path + +import pytest +import torch + +ORCH = Path(__file__).resolve().parent +if str(ORCH) not in sys.path: + sys.path.insert(0, str(ORCH)) + +import test_wiring_e2e as E # noqa: E402 + +from aorlta.config import config_from_dict # noqa: E402 +from aorlta.orchestrator import wiring # noqa: E402 +from aorlta.rollout.compile_check import FakeCompiler # noqa: E402 +from aorlta.rollout.engine_proc import InProcessEngine # noqa: E402 +from aorlta.rollout.fake_engine import BYTE_TOKEN_BASE # noqa: E402 +from aorlta.rollout.templates import ChatFormat # noqa: E402 +from aorlta.sync.writer import expected_count, read_base_index # noqa: E402 + +pytestmark = pytest.mark.heavy + +VOCAB = 151936 # Qwen3's: the byte ids (10000+) and the control ids are real rows +SPECIAL = {"<|endoftext|>": 151643, "<|im_start|>": 151644, "<|im_end|>": 151645, "": 151667, + "": 151668} +_SPECIAL_RE = re.compile("|".join(re.escape(s) for s in SPECIAL)) +_SPECIAL_TEXT = {v: k for k, v in SPECIAL.items()} +IM_END = SPECIAL["<|im_end|>"] + +dataset = E.dataset + + +class Qwen3ByteTokenizer: + """The fake engine's byte tokenizer plus Qwen3's control tokens, with the system/user part of Qwen3's chat + template: thinking off closes an empty think block, thinking on leaves the block to the model. Its format is + read from it as from a real tokenizer (`ChatFormat.derive`).""" + + def __init__(self): + self.chat_format = ChatFormat.derive(self) + + def encode(self, text: str) -> list[int]: + out, i = [], 0 + for m in _SPECIAL_RE.finditer(text): + out += [BYTE_TOKEN_BASE + b for b in text[i:m.start()].encode("utf-8")] + out.append(SPECIAL[m.group(0)]) + i = m.end() + return out + [BYTE_TOKEN_BASE + b for b in text[i:].encode("utf-8")] + + def decode(self, ids) -> str: + parts, buf = [], bytearray() + for t in ids: + if t in _SPECIAL_TEXT: + parts += [buf.decode("utf-8", errors="replace"), _SPECIAL_TEXT[t]] + buf.clear() + else: + buf.append(t - BYTE_TOKEN_BASE) + return "".join(parts + [buf.decode("utf-8", errors="replace")]) + + def render(self, messages, thinking, add_generation_prompt=True) -> str: + text = "".join(f"<|im_start|>{m['role']}\n{m['content'].strip()}<|im_end|>\n" for m in messages) + if not add_generation_prompt: + return text + return text + "<|im_start|>assistant\n" + ("\n\n\n\n" if thinking == "off" else "") + + +def load_engine_model(base_dir) -> torch.nn.Module: + from transformers import AutoModelForCausalLM + + m = AutoModelForCausalLM.from_pretrained(str(base_dir), dtype=torch.bfloat16) + m.eval() + for p in m.parameters(): + p.requires_grad_(False) + return m + + +class RecordingTwin(E.TwinEngine): + """The twin, recording each slot it reloads: its tensors' names, its lm_head.weight if it holds one, and the + reply's count.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.slots: list[dict] = [] + + def update_weights_from_disk(self, model_path: str, flush_cache: bool = True): + from aorlta.sync.writer import load_slot_tensors + + held = load_slot_tensors(model_path) + reply = super().update_weights_from_disk(model_path, flush_cache) + self.slots.append({"names": set(held), "lm_head": held.get("lm_head.weight"), "reply": reply}) + return reply + + +@pytest.fixture(scope="module") +def tiny_qwen3(tmp_path_factory) -> Path: + from safetensors.torch import load_file, save_file + from transformers import Qwen3Config, Qwen3ForCausalLM + + d = tmp_path_factory.mktemp("tiny-qwen3-real-vocab") + torch.manual_seed(1234) + cfg = Qwen3Config(hidden_size=64, intermediate_size=96, num_hidden_layers=2, num_attention_heads=2, + num_key_value_heads=1, head_dim=32, vocab_size=VOCAB, max_position_embeddings=4096, + tie_word_embeddings=True, rms_norm_eps=1e-6) + Qwen3ForCausalLM(cfg).to(torch.bfloat16).save_pretrained(str(d), safe_serialization=True) + path = d / "model.safetensors" + tensors = load_file(str(path)) + assert "lm_head.weight" not in tensors and not (d / "model.safetensors.index.json").exists() + tensors["lm_head.weight"] = tensors["model.embed_tokens.weight"].clone() + save_file(tensors, str(path), metadata={"format": "pt"}) + return d + + +def run_config(tmp_path, base, ds, *, partial: bool, thinking: str): + return config_from_dict({ + "run": {"run_id": "wired-qwen3", "source": "standin", "root": str(tmp_path / "runs")}, + "model": {"base_dir": str(base), "source": "standin"}, + "data": {"prompt_pools": [{"path": str(ds / "listings"), "source": "standin"}]}, + "rollout": {"engines": [0, 1], "K": 4, "P": 4, "max_new_tokens": 512, "request_timeout_s": 300, + "compile_adapter": "fake.compile", "thinking": thinking}, + "reward": {"profile": "fake", "reward_fn": "standin.v11"}, + "orchestrator": {"iterations": 2, "run_seed": 7}, + "trainer": {"ranks": 1, "lr": 1e-2, "micro_batch_tokens": 2048, "gap_gate": "enforce"}, + "sync": {"dir": str(tmp_path / "shm"), "partial": partial}, + }) + + +@pytest.mark.parametrize("partial,thinking", [(True, "off"), (False, "on")]) +def test_two_wired_iterations_on_a_tiny_qwen3_train_publish_and_land(tmp_path, tiny_qwen3, dataset, partial, + thinking): + cfg = run_config(tmp_path, tiny_qwen3, dataset, partial=partial, thinking=thinking) + engines = [RecordingTwin(tiny_qwen3, E.FAKE, i, loader=load_engine_model) for i in range(2)] + handles = [InProcessEngine(i, e) for i, e in enumerate(engines)] + parts = wiring.Parts(tokenizer=Qwen3ByteTokenizer(), handles=handles, compiler=FakeCompiler(), trainer_device="cpu", + probes=E.PROBES) # no load_model: the real loader, and the config's require_fast_path + orch = wiring.build(cfg, parts=parts) + result = orch.run() + assert result["status"] == "completed", result + assert result["iterations_completed"] == 2 + + keys = list(read_base_index(tiny_qwen3)) + assert len(keys) == 25 and "lm_head.weight" in keys + trained_multi_turn = 0 + for n in (1, 2): + s = E.summary(cfg, n) + assert s["samples"]["total"] == 16 and s["scored_fraction"] >= 0.9, s["samples"] + assert s["trainer"]["step_taken"] and s["samples"]["trained"] > 0 + assert s["trainer"]["gap_gate"]["mode"] == "enforce" and s["trainer"]["gap_p99"] < 0.5 + assert s["sync"]["landing"] == "LANDED" and s["sync"]["changed"] > 0 + eps = E.episodes(cfg, n) + trained_multi_turn += sum(1 for e in eps if e["trained"] and len(e["turns"]) > 1) + assert trained_multi_turn > 0 + + merged = orch.trainer.merged + targets = {k for k in keys if re.fullmatch(r"model\.layers\.\d\.(self_attn\.[qkvo]|mlp\.(gate|up|down))_proj\.weight", k)} + assert sorted(merged.weights) == sorted(targets) and len(targets) == 14 # 2 layers x 7 projections + assert all(E.serves(e, merged.weights) for e in engines) + assert not any(torch.equal(merged.weights[k], merged.base[k]) for k in merged.weights) + # the trainer's base was loaded by the real loader, as a causal LM, with the fast-path check required + assert type(orch.trainer.model.get_base_model()).__name__ == "Qwen3ForCausalLM" + # every request stopped on Qwen3's <|im_end|> + reqs = [r for e in engines for r in e.fake.requests if r.outcome == "ok"] + assert reqs and all(r.sampling_params["stop_token_ids"] == [IM_END] for r in reqs) + # each published slot holds what its reload must count: the targets, or every base key, lm_head.weight unchanged + from safetensors.torch import load_file + + base_head = load_file(str(tiny_qwen3 / "model.safetensors"))["lm_head.weight"] + published = engines[0].slots[2:] # after the round trip's two + assert len(published) == 2 + for rec in published: + assert rec["names"] == (targets if partial else set(keys)) + n = expected_count(keys, partial=partial) + assert n == (14 if partial else 25) and rec["reply"][1].endswith(f"({n} checkpoint tensors read)") + assert (rec["lm_head"] is None) if partial else torch.equal(rec["lm_head"], base_head) From c1b915ed83dbc69a9d30b2d5478977d7552550ea Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Fri, 13 Nov 2026 11:11:09 +0000 Subject: [PATCH 3/4] test(aorlta): the Qwen3 end-to-end test's thinking-on replies close their think block, as Qwen3's do Qwen3's thinking-on generation prompt leaves the think block to the model, and the parse reads code only after ; the fake engine's replies had none, so no reply parsed and no repair turn was ever trained on. --- aorlta/tests/orchestrator/test_wiring_e2e_qwen3.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/aorlta/tests/orchestrator/test_wiring_e2e_qwen3.py b/aorlta/tests/orchestrator/test_wiring_e2e_qwen3.py index 71140fb29..c19c3032c 100644 --- a/aorlta/tests/orchestrator/test_wiring_e2e_qwen3.py +++ b/aorlta/tests/orchestrator/test_wiring_e2e_qwen3.py @@ -141,7 +141,13 @@ def run_config(tmp_path, base, ds, *, partial: bool, thinking: str): @pytest.mark.parametrize("partial,thinking", [(True, "off"), (False, "on")]) def test_two_wired_iterations_on_a_tiny_qwen3_train_publish_and_land(tmp_path, tiny_qwen3, dataset, partial, - thinking): + thinking, monkeypatch): + if thinking == "on": # Qwen3's thinking-on prompt leaves the think block to the model: a reply opens and closes it + from aorlta.rollout.fake_engine import FakeEngine + + answer = FakeEngine._answer_text + monkeypatch.setattr(FakeEngine, "_answer_text", + lambda self, rng: "\nshort\n\n\n" + answer(self, rng)) cfg = run_config(tmp_path, tiny_qwen3, dataset, partial=partial, thinking=thinking) engines = [RecordingTwin(tiny_qwen3, E.FAKE, i, loader=load_engine_model) for i in range(2)] handles = [InProcessEngine(i, e) for i, e in enumerate(engines)] From 3f313eba03aef78e2952a34e92b2caa02dc64837 Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Fri, 13 Nov 2026 11:25:48 +0000 Subject: [PATCH 4/4] docs(aorlta): Qwen3-8B configs and what was measured on a GPU; no Qwen3-0.6B config On one MI350X GPU with the arithmetic example: Qwen3-8B scored 256 of 256 samples, 226 correct, landing PROVEN; a 3-iteration training run trained and landed iteration 1, then the gap gate stopped it at iteration 2 (all 8 samples of one prompt 17 to 39 nats from the trainer), not yet explained. Qwen3-0.6B makes TokenSpeed at 3b684074 fail with an illegal memory access after two batches, so its configs are replaced by Qwen3-8B ones. --- aorlta/README.md | 9 ++++++++- aorlta/VALIDATION.md | 7 ++++--- ...h-score-qwen3-0.6b.yaml => arith-score-qwen3-8b.yaml} | 6 +++--- ...h-train-qwen3-0.6b.yaml => arith-train-qwen3-8b.yaml} | 6 +++--- 4 files changed, 18 insertions(+), 10 deletions(-) rename aorlta/configs/{arith-score-qwen3-0.6b.yaml => arith-score-qwen3-8b.yaml} (81%) rename aorlta/configs/{arith-train-qwen3-0.6b.yaml => arith-train-qwen3-8b.yaml} (82%) diff --git a/aorlta/README.md b/aorlta/README.md index fc55f2e82..730aa8c43 100644 --- a/aorlta/README.md +++ b/aorlta/README.md @@ -79,7 +79,14 @@ labels (Section 6.x, C1-C9, D7, ...) of our internal design notes, which are not job mounts from the host, needs ROCm 7.0.2). On a node without the TokenSpeed image the job builds it (about 12 min) from a source tree of that commit, `TS_SRC`. - The base model `Qwen/Qwen3.5-4B` at revision `851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a`, already in the Hugging - Face cache: the job runs offline. + Face cache: the job runs offline. A dense Qwen3 runs too (`src/aorlta/model_family.py`): the control-token ids and + generation prompts come from the model's tokenizer and chat template, and the trainer's loader and the LoRA and merge + names from its `config.json`. Measured on one MI350X GPU, with the arithmetic example: Qwen3-8B + (`configs/arith-score-qwen3-8b.yaml`) scored 256 of 256 samples, 226 correct, landing round trip PROVEN; a + 3-iteration training run trained iteration 1 and its update LANDED, then the gap gate stopped it at iteration 2, + with all 8 samples of one prompt 17 to 39 nats from the trainer's log-probs, which is not explained yet. Qwen3-0.6B + makes TokenSpeed (3b684074) fail with an illegal memory access in its scheduler after two batches, in a score and + a training run alike, so no 0.6B config is shipped. - What the stand-in checks call: aorta at `d22f6116235c05395a0241a83d6f7e952a988454` (a commit of this repository), for its sanitizer backends and test kernels; rocjitsu's `rj_waitcheck` and `librocjitsu_dbi_hooks.so`, from ROCm/rocm-systems branch `shared/rocjitsu/sanitizers`; the ROCm 7.0.2 ASan runtime libraries; a libstdc++ recent diff --git a/aorlta/VALIDATION.md b/aorlta/VALIDATION.md index 49a6f4a78..79394fd5f 100644 --- a/aorlta/VALIDATION.md +++ b/aorlta/VALIDATION.md @@ -122,9 +122,10 @@ was already fair at; it says nothing yet about a real vertical. chosen because it is checkable and has room to improve; the stand-in kernel prompts exercise the loop, not a task. - **The examples' checks are cheap stand-ins** for the checks a real vertical would run (a full test suite, a sanitizer on a real workload, a simulator). -- **Only Qwen3.5-4B, only gfx950, only one node.** Other Qwen3.5 sizes should work and have not been run; Qwen3 - models (such as Qwen3-0.6B) need the changes listed under "Models" in GETTING-STARTED.md; fewer than 8 GPUs has run - only in the 1-GPU arithmetic and Q&A layouts; MI300X and MI355X have not been run. +- **Only Qwen3.5-4B trained, only gfx950, only one node.** Other Qwen3.5 sizes should work and have not been run; a + dense Qwen3-8B scores (the README's base-model entry) but its short training run stopped at iteration 2 on the gap + gate, and Qwen3-0.6B fails in TokenSpeed; fewer than 8 GPUs has run only in the 1-GPU arithmetic and Q&A layouts; + MI300X and MI355X have not been run. - **Long runs**: the longest training run here is 40 iterations. Stability over hundreds of iterations, and resuming one on real hardware beyond the kill-and-resume drill, are not shown. - **Multi-turn episodes beyond compile-repair**, an evaluation exporter during a run, and runs outside Slurm. diff --git a/aorlta/configs/arith-score-qwen3-0.6b.yaml b/aorlta/configs/arith-score-qwen3-8b.yaml similarity index 81% rename from aorlta/configs/arith-score-qwen3-0.6b.yaml rename to aorlta/configs/arith-score-qwen3-8b.yaml index 7b718d0f7..34aa88405 100644 --- a/aorlta/configs/arith-score-qwen3-0.6b.yaml +++ b/aorlta/configs/arith-score-qwen3-8b.yaml @@ -1,15 +1,15 @@ -# arith-score.yaml on Qwen/Qwen3-0.6B (a dense Qwen3, tied embeddings): the 64 held-out multiplications, 4 samples each. +# arith-score.yaml on Qwen/Qwen3-8B (a dense Qwen3): the 64 held-out multiplications, 4 samples each. # One GPU: one TokenSpeed engine on GPU 0 and, for training, one LoRA rank beside it; the answer check runs on the CPU. # Site-specific paths, to edit here or to set in a site file (AORLTA_SITE): run.root, model.base_dir (and # model.adapter_dir to score a trained adapter) and data.prompt_pools[0].path, which is this checkout's # aorlta/examples/qa/arith/heldout. plugins: [examples/qa/qa_task.py] run: - run_id: arith-score-qwen3-0.6b + run_id: arith-score-qwen3-8b source: standin root: /scratch/aorlta/runs model: - base_dir: /data/hf/hub/models--Qwen--Qwen3-0.6B/snapshots/c1899de289a04d12100db370d81485cdf75e47ca + base_dir: /data/hf/hub/models--Qwen--Qwen3-8B/snapshots/b968826d9c46dd6066d109eabc6255188de91218 source: standin data: prompt_pools: diff --git a/aorlta/configs/arith-train-qwen3-0.6b.yaml b/aorlta/configs/arith-train-qwen3-8b.yaml similarity index 82% rename from aorlta/configs/arith-train-qwen3-0.6b.yaml rename to aorlta/configs/arith-train-qwen3-8b.yaml index d4271de34..f193822dd 100644 --- a/aorlta/configs/arith-train-qwen3-0.6b.yaml +++ b/aorlta/configs/arith-train-qwen3-8b.yaml @@ -1,4 +1,4 @@ -# arith-train.yaml on Qwen/Qwen3-0.6B (a dense Qwen3, tied embeddings): 512 generated 3-digit multiplications, 16 per +# arith-train.yaml on Qwen/Qwen3-8B (a dense Qwen3): 512 generated 3-digit multiplications, 16 per # iteration x 8 samples, 20 iterations. # One GPU: one TokenSpeed engine on GPU 0 and, for training, one LoRA rank beside it; the answer check runs on the CPU. # Site-specific paths, to edit here or to set in a site file (AORLTA_SITE): run.root, model.base_dir (and @@ -6,11 +6,11 @@ # aorlta/examples/qa/arith/train. plugins: [examples/qa/qa_task.py] run: - run_id: arith-train-qwen3-0.6b + run_id: arith-train-qwen3-8b source: standin root: /scratch/aorlta/runs model: - base_dir: /data/hf/hub/models--Qwen--Qwen3-0.6B/snapshots/c1899de289a04d12100db370d81485cdf75e47ca + base_dir: /data/hf/hub/models--Qwen--Qwen3-8B/snapshots/b968826d9c46dd6066d109eabc6255188de91218 source: standin data: prompt_pools: