Skip to content

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

Draft
vivekkhandelwal1 wants to merge 5 commits into
feat/aorlta-example-hip-kernelsfrom
feat/aorlta-qwen3-models
Draft

vivekkhandelwal1 wants to merge 5 commits into
feat/aorlta-example-hip-kernelsfrom
feat/aorlta-qwen3-models

Conversation

@vivekkhandelwal1

Copy link
Copy Markdown
Collaborator

Dense Qwen3 models besides Qwen3.5, read from the model rather than hard-coded:

  • Control-token ids and generation prompts come from the model's tokenizer and chat template (templates.ChatFormat): the rollout's stop token, the repair suffix check, the engines' self-test and the replies' stop stripping. A parse adapter no longer names the stop token.
  • The trainer loads the base with its family's class (model_family.py, from config.json). The linear-attention fast-path check doesn't apply to a full-attention model. The LoRA targets, the merge tool and the slot writer accept Qwen3's model.layers.N names, including Qwen3-0.6B's single-file layout with a stored lm_head.weight.
  • A CPU end-to-end test runs two wired iterations on a tiny Qwen3 (real control ids, its chat template, tied embeddings), with thinking off and on.

Measured on one MI350X GPU with the arithmetic example, through the unchanged run job:

  • Qwen3-8B score (configs/arith-score-qwen3-8b.yaml): 256 of 256 scored, 226 correct, landing round trip PROVEN.
  • Qwen3-8B, 3-iteration training: iteration 1 trained (127 of 128 samples) and its update LANDED, with every trainer check passing. The gap gate then stopped the run at iteration 2: all 8 samples of one prompt were 17 to 39 nats from the trainer's log-probs over 20-60% of their tokens, while the other 120 stayed under 0.5. That's systematic rather than numerical noise, and not explained yet; the gate did its job and nothing off-policy was trained on.
  • Qwen3-0.6B: TokenSpeed at 3b684074 fails with an illegal memory access in its scheduler after two clean batches, in a score run and a training run alike, so its configs are replaced by Qwen3-8B ones.
  • Qwen3.5-4B on this branch (regression check, configs/arith-score.yaml): 256 of 256 scored, 243 correct, landing PROVEN.

Stacked on #613. Draft until the iteration-2 gap is understood.

…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.
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.
…heir 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
</think>; the fake engine's replies had none, so no reply parsed and no repair turn was ever trained on.
…n3-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.
Copilot AI balanced review requested due to automatic review settings November 13, 2026 11:26

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 Needs a closer look

Unexplained 17–39-nat Qwen3-8B trainer/engine log-probability divergence at iteration 2 requires GPU/runtime investigation, since CPU twin tests cannot establish TokenSpeed GPU parity.

No open issues 🎉

What changed in this PR

Adds dense Qwen3 support alongside Qwen3.5 in AORLTA’s rollout, training, and weight-sync pipeline, deriving model-specific behavior from the tokenizer and checkpoint configuration. Qwen3-8B scoring is demonstrated; training remains experimental because of the documented iteration-2 log-probability gap.

Changes:

  • Derives control-token IDs and generation prompts from the model’s tokenizer and chat template.
  • Supports family-specific model loading, attention checks, and checkpoint naming.
  • Adds CPU end-to-end coverage, Qwen3-8B arithmetic configs, and documented validation limits.
File Description
aorlta/​VALIDATION.md Records Qwen3 validation results and limitations.
aorlta/​tools/​merge_base_adapter.py Accepts either supported decoder prefix during adapter export.
aorlta/​tests/​trainer/​test_kernels.py Tests when linear-attention checks apply.
aorlta/​tests/​orchestrator/​test_wiring_e2e.py Allows an injectable engine-model loader.
aorlta/​tests/​orchestrator/​test_wiring_e2e_qwen3.py Adds two-iteration Qwen3 CPU integration coverage.
aorlta/​src/​aorlta/​trainer/​rank_pool.py Selects the model loader by checkpoint family.
aorlta/​src/​aorlta/​trainer/​kernels.py Bypasses linear-attention checks for verified full-attention models.
aorlta/​src/​aorlta/​sync/​writer.py Rejects mixed decoder prefixes and documents Qwen3 layouts.
aorlta/​src/​aorlta/​sync/​reload.py Clarifies model-dependent reload counts in documentation.
aorlta/​src/​aorlta/​sync/​merge.py Extends LoRA target-key matching to Qwen3.
aorlta/​src/​aorlta/​rollout/​templates.py Derives chat control IDs and thinking prefixes from tokenizers.
aorlta/​src/​aorlta/​rollout/​feedback.py Uses model-specific repair suffixes for token bounds.
aorlta/​src/​aorlta/​rollout/​fake_engine.py Honors the request’s stop token.
aorlta/​src/​aorlta/​rollout/​engine_pool.py Passes model-specific stop tokens to self-tests.
aorlta/​src/​aorlta/​rollout/​chain.py Uses model stop tokens for generation and reply stripping.
aorlta/​src/​aorlta/​orchestrator/​wiring.py Wires tokenizer-derived stops and validates trainer model families.
aorlta/​src/​aorlta/​model_family.py Defines and validates supported checkpoint families.
aorlta/​README.md Documents Qwen3 support and measured limitations.
aorlta/​examples/​qa/​qa_task.py Removes parser-specific hard-coded stop tokens.
aorlta/​examples/​hip-kernels/​hip_kernels.py Removes parser-specific hard-coded stop tokens.
aorlta/​configs/​arith-train-qwen3-8b.yaml Adds a single-GPU Qwen3-8B arithmetic training config.
aorlta/​configs/​arith-score-qwen3-8b.yaml Adds a Qwen3-8B held-out arithmetic scoring config.

Review effort: Balanced

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants