Repository navigation
Add Qwen 3-8 flash multimodal support - #3916
Merged
Merged
Conversation
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
hubert-marek
changed the base branch from
main
to
chore/remove-transformers-modeling
October 7, 2026 12:17
💡 Codex Review Accept dynamic RoPE checkpoints
When an otherwise supported Llama or Qwen checkpoint declares Invalidate stale Qwen conversion caches
When a user has already trained the composite Qwen3.8 checkpoint, its existing Accept the indexed-attention layer spelling
When a Qwen3.8 checkpoint has been loaded and re-saved through the native Transformers ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
hubert-marek
force-pushed
the
qwen-3-8-flash
branch
from
October 7, 2026 15:38
cd300e0 to
262bd76
Compare
hubert-marek
changed the base branch from
chore/remove-transformers-modeling
to
main
October 7, 2026 15:38
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 262bd76eef
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
hubert-marek
requested review from
garrett361 and
samsja
and removed request for
S1ro1
October 9, 2026 15:44
This was referenced Oct 10, 2026
samsja
pushed a commit
that referenced
this pull request
Oct 19, 2026
Qwen3.8's ngram table uses vocabulary-row shards, but deployed PyTorch 2.13 DTensor dense embedding backward allocates the full table gradient on each rank. The 320,001,536 × 160 table would require a **95.37 GiB BF16 gradient allocation per rank**. Keep the original row-sharded weights and head mesh, and perform `F.embedding` on plain local weights so backward allocates only local rows. Route each flattened ID directly to the rank that owns it instead of gathering everyone's IDs and computing masked lookups on every rank: 1. Sort local IDs to group them by row owner and exchange per-peer counts. 2. Stack send/receive counts and copy both to CPU once; the stack consumes the asynchronous count collective before its values become split lists. 3. Send IDs to owners with variable-split all-to-all, and look up only received local rows. 4. Return vectors with autograd-enabled all-to-all and restore local token/head order with the saved permutation. Its backward sends vector gradients back to the owners. No token padding, clamping, non-owner masking or embedding-output reduce-scatter is needed. Uneven vocabulary shards, empty row shards, uneven/empty token batches and arbitrary input shapes are supported. Sparse, frequency-scaled, max-norm and `padding_idx` embeddings fail explicitly before partitioning; the Qwen ngram table uses none of these modes. Current `main` already routes ngram embeddings and expert biases to AdamW through #4108. This PR retains that upstream optimizer behavior and changes only the embedding lookup implementation. Targets `main`, which already contains native text-only Qwen 3.8 support. The diff is limited to `embedding_parallel.py`; multimodal support remains in #3916. Weight placement, parameter names and row-shard checkpoint layout remain unchanged. Companion #4102 bounds reconstruction memory during ngram weight broadcast. This PR addresses lookup/backward allocations; it does not avoid full-table gathering in weight transfer by itself. Merged current `main` on 2026-10-16 to resolve the optimizer conflict with #4108. The embedding implementation is unchanged (same tested SHA-256 below). ### Latest validation Allocated India H200 compute, job **2875**, deployed **PyTorch 2.13.0+cu130**. All compute ran on the allocated node, not the login node. NCCL NVLS was disabled for the test after a separate attempt failed during communicator initialization before lookup. - **NCCL at 1, 2, 4 and 8 GPUs: 375 passing rank-case records.** FP32/BF16 output and full-gradient parity against unsharded `nn.Embedding`, equal/uneven token counts, one/all empty ranks, vocabulary sizes not divisible by world size and empty row shards. Each rank runs three AdamW updates with two microbatches, including activation-checkpoint recomputation; updates match within FP32 tolerance. - Instrumented lookup counts match exactly the number of IDs owned by that rank. Dense backward receives only its local row count and returns a plain local gradient. These checks also exercise completion of the asynchronous count exchange before split lists are used. - **Four-GPU FSDP:** BF16 computation over FP32 parameters/reduction; outputs and full gradients match exactly across equal/uneven/empty batches. AdamW updates and activation checkpointing pass. Distributed model and optimizer checkpoint restore is exact, a subsequent update matches the reference, and the saved weights load exactly through the original rowwise partitioner. - Additional four-GPU checks pass for a zero-row vocabulary with empty input, int32 IDs, noncontiguous 3D inputs, 1D inputs and explicit rejection of unsupported embedding modes. - The unchanged Muon-to-AdamW exclusion was previously checked against the baseline optimizer constructor on 12 parameter fixtures: trainable ngram tables move to AdamW, other Muon/AdamW memberships and frozen exclusions stay unchanged. - Ruff, formatting and `git diff --check` pass. The tested module SHA-256 is `d26427312f5e3d5c2aba43dc7f03722e03e5c7f0f997918b6662075838ff9d4d`. A separate four-GPU 65,536 × 1,024 FP32 table check confirms that DTensor rowwise dense backward allocates **256 MiB**, while this implementation's local backward allocates **64 MiB** per rank. For Qwen's actual table at 16 ranks, the local BF16 gradient is approximately **5.96 GiB per rank**, calculated from dimensions rather than measured full-model memory. ### Lookup memory and time Matched synthetic forward/backward benchmark: BF16 65,536 × 160 table, **16,384 local tokens × 16 heads**, balanced row ownership, one H200 node, NCCL with NVLS disabled. Compared the previous all-gather/mask/reduce-scatter implementation with the new owner-routed implementation. Three warmups and ten timed iterations; wall time includes count exchange and host synchronization. Reported values are the maximum across ranks. Peak extra allocation includes lookup temporaries, output and local dense gradient, excluding preallocated weights and IDs. | GPUs | Previous peak MiB | Owner-routed peak MiB | Previous forward/backward ms | Owner-routed forward/backward ms | |---|---:|---:|---:|---:| | 2 | 331.0 | 248.0 | 5.76 | 2.44 | | 4 | 661.0 | 249.0 | 13.71 | 4.45 | | 8 | 1318.0 | 249.0 | 29.19 | 4.42 | With balanced ownership, lookup buffers stay approximately constant as world size increases in this measurement. **This is not a worst-case guarantee:** sending all IDs to one owner raises its peak. At eight GPUs, the deliberately skewed case measured **1,318.0 → 842.3 MiB** peak and **32.36 → 6.84 ms** maximum forward/backward time. At two GPUs, the skewed owner-routed peak was **332.0 MiB**, slightly above the previous **331.0 MiB**. Deduplicating repeated IDs is a possible follow-up and is not included here. These are isolated embedding measurements, not full-model training throughput or memory results. Full-model weight broadcast and compiled-model integration remain unverified. Migration of optimizer checkpoints whose ngram table previously belonged to Muon is not covered; these tests also do not establish the cause of the separate router-replay indexing failure. Validation scripts and results are private under `/data/nvme0/hubert/qwen38-embedding-oom-20261008/local-row-backward/owner-routed-20261013/` and are not included in the diff. No live training run or shared environment was changed. The automated summary below describes an earlier commit. <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **High Risk** > Changes distributed embedding collectives and gradient accounting on the training critical path; incorrect routing or autograd would silently break parity or OOM large ngram shards. > > **Overview** > Replaces vocabulary-parallel embedding forward/backward so **row-owned IDs are all-to-all routed** and lookups run on **plain local weights**, avoiding PyTorch 2.13 DTensor dense backward that would allocate a full-table gradient per rank (critical for large ngram tables). > > The old **token-batch padding + RowwiseParallel prepare hooks** are removed in favor of sort/bincount, variable-split `all_to_all` for IDs and autograd-enabled return of vectors, then permutation restore to the original token order. Unsupported embedding modes (`padding_idx`, sparse, etc.) are rejected at apply time. > > **Optimizer:** `ngram_embedding` parameters are excluded from Muon and trained with the existing AdamW group, same as token embeddings and `lm_head`. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit 7cfdedc. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY --> Co-authored-by: Codex <codex@local>
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 049f2b5f37
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
hubert-marek
added a commit
that referenced
this pull request
Oct 21, 2026
Compiled Qwen 3.8 training fails on a single packed sequence with `ConstraintViolationError`: the ngram/position-learning path specializes `cu_seqlens` to length 2, but `mark_dynamic` forbids that specialization. Use `maybe_mark_dynamic` for this dimension so Dynamo can specialize the single-sequence case while attempting dynamic shapes for other packed-sequence counts. Shape guards remain active. This changes only Qwen 3.8; it does not change embedding sharding, model math, or other architectures. Targets `main`, which already contains the affected text-model path. Multimodal support in #3916 is independent of this fix. Validation on an H200 compute allocation with PyTorch `2.13.0+cu130`: - The unmodified strict annotation reproduces the same `ConstraintViolationError` in the real ngram module. - The optional annotation passes 10 compiled/eager PLE comparisons: one, two, four, one and three packed sequences, both with and without non-reentrant activation checkpointing. Outputs, input/parameter gradients and AdamW parameter updates match within `rtol=2e-4, atol=2e-5`. - Ruff check, formatting check and `git diff --check` pass. The isolated probe uses small real PLE/ngram modules, not the full distributed model. These checks do not establish a successful full-model distributed compiled update. <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **Low Risk** > Single-line Dynamo annotation change in one model; validated for numerical parity across packed-sequence shapes without altering training math. > > **Overview** > **Fixes compiled Qwen 3.8 training** when a batch is a single packed sequence: Torch Dynamo was hitting `ConstraintViolationError` because the PLE/ngram path specializes `cu_seqlens` length (e.g. to 2) while `mark_dynamic` blocked that specialization. > > In `Qwen3_8FlashNextTextModel.forward`, **`cu_seqlens` dim 0** is annotated with `torch._dynamo.maybe_mark_dynamic` instead of `mark_dynamic`, so Dynamo can specialize the one-sequence case while still treating other packed-sequence counts as dynamic. Scope is **Qwen 3.8 Flash Next only**; model math, sharding, and other architectures are unchanged. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit afcb32d. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY -->
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 2 potential issues.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit dee2687. Configure here.
hubert-marek
pushed a commit
that referenced
this pull request
Oct 24, 2026
hubert-marek
added a commit
that referenced
this pull request
Oct 24, 2026
CPU-only orchestrators crash before collecting rollouts because `BatchPacker` builds a Qwen3.8 model on the `meta` device for FLOP counting, while its sparse-attention constructors immediately load GPU-only `indexed_attention` kernels. No attention forward pass is needed for this packing-cost estimate. Guard both constructor-time kernel loads with `torch.cuda.is_available()`. CPU/meta construction can then count model shapes and attention FLOPs; GPU trainers still resolve the same kernels before `torch.compile`. Forward/backward code, weights and numerical settings are unchanged. The existing model checks now cover full-attention layers as well as linear-attention layers. A regression check constructs text/composite models on `meta` with CUDA reported unavailable and evaluates `FlopCounter`. Validation on an H200 with PyTorch 2.13.0+cu130 and Transformers 5.6.2: | Check | Result | | --- | --- | | Upstream model checks, CUDA hidden | 12 passed | | Upstream model checks, CUDA available | 12 passed | | Combined-runtime Qwen3.8/Qwen3.5 model and vision checks | 29 passed | | Real `BatchPacker`, full checkpoint config, 237,568-token cost, CUDA hidden | Passed on upstream and combined runtime; 12 indexed layers and identical FLOP counts | | Real indexed-attention GPU forward/backward | Passed on upstream and combined runtime; kernels bound before forward, finite outputs and nonzero finite input gradients | | Ruff and whitespace checks | Passed | These are targeted startup and kernel/model checks, not a completed full-size training step. The fix is also applied unchanged to the Model Factory runtime alongside #3916 and #4241. <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **Low Risk** > Narrows kernel loading to CUDA hosts; CPU/meta paths are for FLOP counting only, so training forward behavior on GPU is unchanged. > > **Overview** > Fixes CPU-only orchestrator crashes when **BatchPacker** builds Qwen3.8 on the **meta** device for FLOP counting: sparse attention no longer loads GPU-only `prime_kernels` `indexed_attention` at `__init__` unless `torch.cuda.is_available()`. > > **`IndexedGatedAttention`** and **`SparseAttentionIndexer`** still bind the same kernels on GPU trainers (before compile); only construction timing changes—forward/backward and weights are unchanged. > > Tests now parametrize **`linear_attention`** and **`full_attention`**, and add **`test_meta_flop_count_without_cuda`** to assert meta construction and **`FlopCounter`** work with CUDA reported unavailable. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit a25a85b. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY --> Co-authored-by: Codex <codex@local>
hubert-marek
added a commit
that referenced
this pull request
Oct 24, 2026
Adds editable image-RL templates for Qwen3.5-122B-A10B and Qwen3.8-Flash-Next, plus README, model-support and multimodal documentation. The recipes preserve our 4-trainer/4-inference topology, 237,568-token context, 184,320-token browser compaction, Muon/IPO recipe, frozen vision encoder, inference precision/caching, and shuffled sampling. Batch/group sizes are 336/16 and 256/16 respectively. Both examples use an illustrative Muon learning rate of `1e-6` and `max_steps = 200`. Adaptive episode concurrency starts at 96, with a ceiling of 512 rather than a fixed target. Environment/harness IDs, run names, W&B identity and shared paths use explicit placeholders. The concise README covers customization and launch, with links to the full setup, scaling and training docs. Browser-specific reward/compaction fields are called out explicitly. No private catalogs, run IDs or secrets are included. Also corrects the outdated claim that VLMs require BF16 optimizer/reduction dtypes; the templates retain FP32. Validation: both TOMLs pass the actual RLConfig schema with only environment/harness IDs substituted for locally available browser plugins, component-config serialization, JSON round trips and checks of numerical/topology/sampling settings. All 66 relative documentation links and whitespace checks pass. No new GPU training run was launched for this documentation change. Stacked on #3916 (`qwen-3-8-flash`), so this diff contains only examples and documentation. Merge after #3916 and retarget to main; the broadcast/reload and CPU/meta fixes (#4102, #4241, #4246) are already merged into main. --------- Co-authored-by: Codex <codex@local>
Co-authored-by: Cursor <cursoragent@cursor.com>
Adds editable image-RL templates for Qwen3.5-122B-A10B and Qwen3.8-Flash-Next, plus README, model-support and multimodal documentation. The recipes preserve our 4-trainer/4-inference topology, 237,568-token context, 184,320-token browser compaction, Muon/IPO recipe, frozen vision encoder, inference precision/caching, and shuffled sampling. Batch/group sizes are 336/16 and 256/16 respectively. Both examples use an illustrative Muon learning rate of `1e-6` and `max_steps = 200`. Adaptive episode concurrency starts at 96, with a ceiling of 512 rather than a fixed target. Environment/harness IDs, run names, W&B identity and shared paths use explicit placeholders. The concise README covers customization and launch, with links to the full setup, scaling and training docs. Browser-specific reward/compaction fields are called out explicitly. No private catalogs, run IDs or secrets are included. Also corrects the outdated claim that VLMs require BF16 optimizer/reduction dtypes; the templates retain FP32. Validation: both TOMLs pass the actual RLConfig schema with only environment/harness IDs substituted for locally available browser plugins, component-config serialization, JSON round trips and checks of numerical/topology/sampling settings. All 66 relative documentation links and whitespace checks pass. No new GPU training run was launched for this documentation change. Stacked on #3916 (`qwen-3-8-flash`), so this diff contains only examples and documentation. Merge after #3916 and retarget to main; the broadcast/reload and CPU/meta fixes (#4102, #4241, #4246) are already merged into main. --------- Co-authored-by: Codex <codex@local>
Use linear patch projection, differentiable learned-position interpolation and activation-dtype rotary trig to align vision arithmetic with inference. Preserve checkpoint layouts and native autograd; expand existing VLM backward coverage to fractional grids.
hubert-marek
force-pushed
the
qwen-3-8-flash
branch
from
October 26, 2026 15:36
cd0cda3 to
7370c05
Compare
samsja
previously approved these changes
Oct 27, 2026
samsja
approved these changes
Oct 27, 2026
faresobeid
added a commit
that referenced
this pull request
Nov 2, 2026
…points load without a vision tower (#4335) Since #3916 (Qwen3.8) and the Qwen3.5 VLM unification, text-only Qwen3.5 (`qwen3_5`, `qwen3_5_moe`) and Qwen3.8 (`qwen4_exp`) composite checkpoints no longer load in the trainer. These configs have `language_model_only: true` and a `vision_config`, but their weights hold only `model.language_model.*` (plus `lm_head`). **Symptom.** On main, SFT on such a checkpoint fails while loading: `RuntimeError: Missing key in checkpoint state_dict: model.visual.blocks.0.attn.proj...`. Meta construction shows the cause: the model has `model.visual.*` keys that the checkpoint does not contain (153 for Qwen3.5-0.8B with the vision weights removed, 333 for a 4-layer Qwen3.8). **Root cause.** Both composite configs dropped `language_model_only`, and `vision_config` falls back to a default. `ForCausalLM` sets `is_vlm = hasattr(config, "vision_config")`, so it always builds the vision tower. **Fix (same for both models).** Keep `language_model_only` on the config. This follows vLLM's `--language-model-only`; official Qwen configs ship the field. When the flag is set: - the composite model sets `visual = None` and its forward calls the language model directly (the injected LM-head forward calls `model.model(...)`, so the branch lives there); - `is_vlm` and `supports_packed_multimodal_training` are False. The flag is kept on the config, so saved configs keep it. Key names stay `model.language_model.*`, so conversion and the weight names sent to vLLM do not change. `docs/advanced.md` notes that vLLM does not read the field, so text-only RL also needs `inference.vllm.language_model_only = true`. I did not derive that vLLM setting automatically. The config package never reads model files (the model can be a Hub id), so doing it there would mean loading the HF config inside config validation. **Test.** - Qwen3.8: the config fixture of `test_model.py` gets a `language_model_only` case. - Qwen3.5: a meta test checks that a `language_model_only` config has no vision encoder and only `model.language_model.*` / `lm_head` keys. GPU (H200): - Qwen3.5-0.8B with the vision weights removed and `language_model_only: true`: SFT fails on main with the missing-key error. With the fix, SFT runs 8 steps. - Qwen3.8, 4 layers, text-only: SFT on 4 GPUs, 8 steps, loss 12.88 -> 10.98. RL with vLLM `language_model_only = true` (4 trainer / 4 inference GPUs, NCCL broadcast) runs 6 steps, sending 133 layer state dicts per step, mismatch KL 0.075-0.09. - VLM checkpoints (Qwen3.5-0.8B and a tiny Qwen3.8) with `[model.vlm]`: the vision encoder is still built, frozen and wrapped with FSDP. <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **Medium Risk** > Changes Qwen3.5/Qwen3.8 model construction and checkpoint loading paths; behavior is scoped to the `language_model_only` flag with new unit coverage, but mistakes could affect weight loading or VLM vs text-only routing. > > **Overview** > Fixes loading and training for **text-only** Qwen3.5 / Qwen3.8 composite checkpoints that still carry a `vision_config` but only `model.language_model.*` weights (official `language_model_only: true` in `config.json`). > > **Trainer behavior:** The flag is preserved on `Qwen3_5Config` and `Qwen3_8FlashNextConfig` instead of being dropped. Composite `ForCausalLM` models now treat VLM mode as `vision_config` present **and** not `language_model_only`; the composite wrapper skips building `visual`, forwards through the language model only, and clears `is_vlm` / packed multimodal support. `init_buffers_post_meta` uses `get_language_model` so buffer setup still targets the decoder when the wrapper is composite but text-only. > > **Docs:** Notes that vLLM ignores the HF field, so text-only RL must set `inference.vllm.language_model_only = true`. > > **Tests:** Qwen3.5 meta test asserts no vision encoder and state dict keys under `model.language_model.*` / `lm_head`; Qwen3.8 config fixture adds a `language_model_only` variant. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit f30d782. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY -->
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.
Adds native image training to Qwen3.8 Flash Next (
qwen4_exp) through Prime-RL’s existing multimodal pipeline.primeintellect/color-codeword@0.1.0environment and nativenullharness, with disjoint training/evaluation slices and comments showing where to substitute your multimodal environment. Qwen3.8 uses CP8.Rebased onto main
088570f8(2026-10-26). Main now supplies #4252’s PLE memory fix, #4264’s FLA 0.6.0 update and #4279’s new QSA kernels/activation-checkpoint plan handling.Validation
Nebius H200, PyTorch 2.13.0+cu130, exact release kernel wheel
v0.1.0-b52bd9d, job 6812:test_load_configscases pass (2 passed). Both configs also load 1,024 shuffled image training tasks and 128 disjoint evaluation tasks, and resolveColorCodewordEnvwithNullHarness. Browser-only fields were replaced with the native harness's compaction setting. Both examples inherit the default Prime runtime (no subprocess override); train/eval runtime resolution was verified. Whitespace checks pass.Earlier frozen-model validation of the three vision changes used one eight-image prefix and all 4,545 generated targets, with saved inference expert replay and stable attention ordering. Exact full-vocabulary KL fell 0.002570 → 0.001602 with FLA 0.5.2 and 0.002430 → 0.001249 with FLA 0.6.0. These results predate this main rebase and do not establish training-wide gains.
Note
Medium Risk
New vision encoder and MRoPE/CP paths affect large-scale RL training numerics and distributed behavior; mitigated by extensive unit tests but not full production KL validation in this PR.
Overview
Adds native vision-language training for Qwen3.8 Flash Next (
qwen4_exp) on top of Prime-RL’s existing multimodal stack, alongside doc and example updates for Qwen3.5/Qwen3.8 image RL and SFT.The trainer gains a composite VLM model: vision encoder (
model.visual), image-token scatter into language embeddings, packed MRoPE positions frommm_token_type_ids/image_grid_thw, and context-parallel sharding after vision merge. PLE still hashes raw placeholderinput_ids; checkpoint conversion keepsmodel.visualweights instead of dropping them.QwenVLAdapterregistersqwen4_exp.Docs list Qwen3.8 in model/VLM tables, add a
qwen3.8renderer, and revise multimodal limits (FP32 trainer + BF16 vLLM is OK; monitor mismatch KL metrics).Examples: new
rl-qwen3.5.tomlandrl-qwen3.8.toml(4×8 H200 trainer + 4 inference replicas,color-codeword, GRPO/IPO); expandedexamples/extra/vlm/README.md.Tests cover VLM forward/backward, router replay, MRoPE packing, and composite checkpoint roundtrip.
Reviewed by Cursor Bugbot for commit b12f244. Bugbot is set up for automated code reviews on this repo. Configure here.