Repository navigation
feat(aorlta): dense Qwen3 models, from the tokenizer and the checkpoint's config - #616
Draft
vivekkhandelwal1 wants to merge 5 commits into
Draft
vivekkhandelwal1 wants to merge 5 commits into
vivekkhandelwal1 wants to merge 5 commits into
Conversation
…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.
Contributor
There was a problem hiding this comment.
🔵 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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Dense Qwen3 models besides Qwen3.5, read from the model rather than hard-coded:
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.model_family.py, fromconfig.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'smodel.layers.Nnames, including Qwen3-0.6B's single-file layout with a storedlm_head.weight.Measured on one MI350X GPU with the arithmetic example, through the unchanged run job:
configs/arith-score-qwen3-8b.yaml): 256 of 256 scored, 226 correct, landing round trip PROVEN.configs/arith-score.yaml): 256 of 256 scored, 243 correct, landing PROVEN.Stacked on #613. Draft until the iteration-2 gap is understood.