feat(aorlta): dense Qwen3 models, from the tokenizer and the checkpoint's config

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

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.

aorlta/configs/arith-score-qwen3-8b.yaml

@@ -0,0 +1,38 @@
+# 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-8b
+  source: standin
+  root: /scratch/aorlta/runs
+model:
+  base_dir: /data/hf/hub/models--Qwen--Qwen3-8B/snapshots/b968826d9c46dd6066d109eabc6255188de91218
+  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

aorlta/configs/arith-train-qwen3-8b.yaml

@@ -0,0 +1,43 @@
+# 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
+# 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-8b
+  source: standin
+  root: /scratch/aorlta/runs
+model:
+  base_dir: /data/hf/hub/models--Qwen--Qwen3-8B/snapshots/b968826d9c46dd6066d109eabc6255188de91218
+  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

aorlta/examples/hip-kernels/hip_kernels.py

@@ -97,11 +97,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:"):

aorlta/examples/qa/qa_task.py

@@ -89,11 +89,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
 

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)

aorlta/src/aorlta/orchestrator/wiring.py

@@ -346,13 +346,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)
 
@@ -1279,6 +1281,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

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)

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:

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):