Skip to content

feat(sft): Enable VLM SFT #2485

Closed
eligotts wants to merge 63 commits into
mainfrom
feat/sft-multimodal-renderer
Closed

feat(sft): Enable VLM SFT #2485
eligotts wants to merge 63 commits into
mainfrom
feat/sft-multimodal-renderer

Conversation

@eligotts

@eligotts eligotts commented May 13, 2026

Copy link
Copy Markdown
Contributor

Summary

This PR makes SFT renderer-only and enables VLM SFT end-to-end, and builds out the Qwen3.5 (dense + MoE) custom model stack needed to train those models with context parallelism.

Renderer-only SFT

  • Samples always go through renderers.build_training_sample; the legacy incremental tokenization path is removed. [renderer] defaults to auto (hand-coded renderer required; DefaultRenderer is rejected).
  • deps/renderers advances to 99bedaf, which surfaces the multimodal payload through build_training_sample.

Multimodal SFT plumbing

  • mm_kwargs, mm_token_type_ids, and seq_lens flow through datasets, packing, collation, and trainer forward().
  • A single CatDataset packing path serves text-only and multimodal samples: text-only spans can share a multimodal pack with zero mm_token_type_ids; compatible multimodal kwargs concatenate along the leading item dimension. pack_function = "cat" is the SFT default again (stack remains opt-in).
  • The trainer loads an AutoProcessor, attaches it to the renderer, saves it with VLM weight checkpoints (so AutoProcessor.from_pretrained(step_dir) works on exported dirs), and enforces VLM rules (custom PrimeRL VLM implementations only — the generic HF VLM path is rejected; image-only — video inputs are rejected at sample processing; micro_batch_size = 1, LoRA vs unfrozen vision, fail-fast when no multimodal processor can be loaded). New examples/vlm_sft and examples/vlm_sft_moe configs (data paths are placeholders).
  • data_files cannot be combined with subsets/splits (config-rejected: subsets were silently ignored and non-train splits failed at load time).
  • Local .jsonl/.zst data_files are normalized into a $TMPDIR cache before load_dataset; the cache (and zstd decompression output) is written to a .part sidecar and atomically renamed into place, so an interrupted run can't leave a truncated cache that later runs silently reuse.
  • pack_function = "stack" is rejected at config for Qwen3.5 custom flash attention: its varlen path assumes a single flattened sequence (batch size 1), which stack's genuine multi-row batches violate, crashing with a shape mismatch.
  • Non-flash (sdpa/eager) attention now builds cu_seqlens for text-only cat packs, not just 3D MRoPE packs: linear-attention (GatedDeltaNet) state correctly resets at pack boundaries, and multi-document packs with full_attention layers raise (SDPA can't mask document boundaries) instead of silently letting later samples attend to earlier ones.
  • data.type = "fake" no longer requires (or constructs) a renderer — fake data never renders messages, so benchmarking a model without a hand-coded renderer works again.
  • Glob patterns in data_files (e.g. data/*.jsonl) are expanded before normalization instead of being stat()-ed literally.
  • tool_defs (verifiers rollout format) is JSON-stringified before the Arrow load, same as tools — heterogeneous per-row tool schemas otherwise crash schema inference or gain phantom null fields.
  • Arrow-injected phantom nulls are stripped from messages before tool-call deserialization, so genuine null values inside tool-call argument JSON survive into the rendered training text.
  • data.loss_mask.user/system/tool = true works again under renderer-only SFT via the renderer's body-only path (content_sft_roles): the message content is trained, but not the role scaffolding (e.g. <tool_response> tags), unlike the legacy path which trained the full message span — loss for such configs is not bit-identical to pre-renderer runs.

Qwen3.5 custom stack + context parallelism

  • New dense Qwen3_5ForCausalLM custom implementation (HF vision tower + PrimeRL text stack), so impl = "custom" routes dense Qwen3.5 full-attention layers through PrimeRL's patchable attention classes (FA2/FA3/FA4, ring/ulysses CP) — same as Qwen3.5 MoE.
  • Hybrid CP for the GatedDeltaNet (linear-attention) layers via fla's native CP ops: hook-only setup through set_context_parallel_attributes, real global document boundaries forwarded via seq_lens (seq_lens_are_global), CP-aware conv1d, and cross-rank recurrent-state exchange.
  • The hook is the single CP-setup protocol: setup_model_cp replaces setup_hybrid_cp + setup_nemotron_h_cp in both trainers, NemotronH exposes the hook at the top level (fanning out to its Mamba layers), softmax-only models no-op, and only a hybrid model without the hook raises.
  • CP context setup is sync-free: the global-vs-local cu_seqlens decision is threaded down as a provenance bool instead of reading the CUDA tensor back per linear-attention layer per microbatch.
  • VLM CP (ulysses): image batches stay full-length through image scatter/MRoPE, then the model shards sequences internally; the vision encoder stays on SDPA.

CP correctness fixes (root-caused during E2E validation on this branch)

  • Document-boundary packing can produce packs shorter than seq_len with odd lengths. shard_for_cp used torch.chunk, which splits odd lengths unevenly across CP ranks, so the ulysses all-to-all size-mismatched within a CP pair and NCCL deadlocked before step 0 (diagnosed via the NCCL flight recorder: mesh_cp all_to_all input sizes [2, 699, 8, 256] vs [2, 698, 8, 256] on the stuck pair). Fixes:
    • CatDataset pads every pack to exactly seq_len (pad tokens form their own document and are excluded from the loss) — the single padding site; the trainer does not re-pad.
    • shard_for_cp validates divisibility and batch shape up front, so an uneven pack now fails loudly instead of hanging the job.

E2E validation at PR head 71c82e4e0

Setup: 8×H100 (single node), impl = "custom", attn = "flash_attention_3", seq_len = 8192, pack_function = "cat", LoRA, 10 steps on a normalized 10k-row JSONL slice of PrimeIntellect/INTELLECT-3-SFT-10K, W&B project qwen35-custom-attention-e2e. CP rows use CP 2 (ulysses). Resolved settings inherited from current main defaults: compile on, activation offloading on, optim_cpu_offload = true.

run model CP EP loss impl steps losses (steps 0→9) step time peak mem W&B
dense text CP Qwen/Qwen3.5-0.8B 2 1 liger_fused 10/10 0.9662, 0.8727, 0.8720, 0.6853, 0.8664, 0.9094, 0.8324, 0.9404, 0.9453, 0.7762 ~850ms 2.1 GiB pl1ncv8t
MoE text CP Qwen/Qwen3.5-35B-A3B 2 4 liger_fused 10/10 0.6106, 0.5296, 0.5446, 0.4061, 0.5366, 0.5684, 0.5607, 0.5626, 0.6025, 0.4784 ~13.5s 14.7 GiB fa9dtakv
dense text no-CP Qwen/Qwen3.5-0.8B 1 1 liger_fused 10/10 1.0451, 0.7927, 0.8430, 0.7929, 0.8006, 0.9178, 0.8406, 0.9146, 0.9459, 0.7773 ~560ms 2.7 GiB kcqhjnbu
dense text no-CP Qwen/Qwen3.5-0.8B 1 1 torch 10/10 1.0451, 0.7925, 0.8432, 0.7928, 0.8009, 0.9184, 0.8406, 0.9144, 0.9462, 0.7775 ~2.5s 28.7 GiB lc0g7nrc
MoE text no-CP Qwen/Qwen3.5-35B-A3B 1 4 liger_fused 10/10 0.6568, 0.4835, 0.5117, 0.4930, 0.4874, 0.5857, 0.4913, 0.5689, 0.5646, 0.4920 ~10.4s 15.9 GiB swn527ch
MoE text no-CP Qwen/Qwen3.5-35B-A3B 1 4 torch 10/10 0.6570, 0.4842, 0.5115, 0.4926, 0.4875, 0.5851, 0.4913, 0.5686, 0.5649, 0.4919 ~14.7s 36.7 GiB 16czkxsr

Notes:

  • Pre-PR baseline is unrunnable on this dataset: at the merge-base (436873190), the legacy incremental-tokenization path skips 100% of samples with the Qwen3.5 chat template ("Mismatch in incremental tokenization … not stable under incremental application"; 190k+ skip warnings, zero training steps before timeout — lq05gysm, wxtpy2p7). Renderer-only SFT is what makes Qwen3.5 SFT possible at all on such data.
  • torch vs liger_fused parity: dense losses match step-for-step to ≤2e-4, MoE to ≤1e-3 (grouped-GEMM nondeterminism). torch costs much more memory (full logits) and step time.
  • Both CP configurations hung before step 0 prior to the CP padding/validation fixes (see above); the same configs now train to completion with stable losses and nan_count = 0.
  • Losses are stable across heads: the dense CP run reproduces step-for-step (±1e-3) on cfd518c93, 6a01021f4, and 71c82e4e0.
  • CP and no-CP rows are not loss-comparable to each other: dataset sharding happens before packing (4-way with CP vs 8-way without), so pack compositions differ. Within a parallelism setting, the loss-impl pairs match.
  • MoE step time (~10–15s) is dominated by main's new activation-offloading + optimizer-CPU-offload defaults, not by this PR: the identical CP config with those settings off (older defaults) runs at ~2.2s/step, 29k tok/s, MFU 11.9%, 40.9 GiB peak. Dense is slightly faster with the new defaults (compile wins at 0.8B scale).

CI: reverse_text_sft::test_loss_goes_down failure analysis

The failing check is a flakiness issue exposed by the renderer change, not a training regression:

  • Deterministic repro on H100 matches CI exactly (PR head: 5.5987, 5.4090, 5.4729, 5.5753, 5.6005, grad norms 60–80). The pre-PR control on the same hardware shows the same grad-norm profile (65–72) and the same ±0.2–0.3 per-step loss noise (6.0826, 5.9109, 6.0219, 6.2301, 5.9531) — it "passes" only because its step 4 happens to land below step 0. At lr = 1e-6 with max_norm = 1.0 over 5 steps, true optimization movement is far smaller than the inter-batch noise, so the start-vs-end comparison is a coin flip on both revs.
  • What changed with the renderer: the qwen3 renderer includes the assistant's empty <think>\n\n</think> block in the loss mask (legacy masking started at the reversed-text payload), which lowers the absolute loss (~6.0 → ~5.5) and re-rolls pack composition. Target/input alignment and packing were verified identical/correct on both revs.
  • Suggested fix (separate change): make the check robust — average the first-k vs last-k steps, or raise the LR / step count so the learning signal exceeds batch noise.

Prior validation (earlier heads)

Loss parity, packing benchmarks, and earlier E2E matrices (validated on earlier branch heads)

Loss parity: renderer path vs legacy path (dense 0.6B, 20 steps)

  • Dense 0.6B 20-step online W&B H100 SFT 2x2 completed on prime-job-ssh-node-5f2d5979-0, project packing, with Qwen/Qwen3-0.6B, PrimeIntellect/INTELLECT-3-SFT-10K (default / math), batch_size=1, micro_batch_size=1, seq_len=1024, pack_function="cat", shuffle=false, seed=0, and max_steps=20. Before bf3d381 used renderers 1933293; after f9185a7 used renderers b4c4051. Before/after losses match exactly at logged precision for both loss implementations; torch is the repo's non-Liger CE path.
version loss_impl W&B peak mem losses, steps 0-19
before liger_fused sft-mm-before-liger-20 6.0 GiB 1.0363, 1.3208, 0.9778, 0.5132, 1.0865, 0.8884, 1.0739, 0.9028, 0.8801, 0.6700, 0.7327, 0.5879, 0.8547, 0.8070, 0.5961, 0.7454, 0.5147, 1.2197, 0.4154, 0.4856
before torch sft-mm-before-torch-20 6.3 GiB 1.0363, 1.3195, 0.9791, 0.5142, 1.0861, 0.8911, 1.0728, 0.9029, 0.8784, 0.6701, 0.7335, 0.5883, 0.8535, 0.8060, 0.5962, 0.7453, 0.5149, 1.2208, 0.4153, 0.4853
after liger_fused sft-mm-after-liger-20 6.0 GiB 1.0363, 1.3208, 0.9778, 0.5132, 1.0865, 0.8884, 1.0739, 0.9028, 0.8801, 0.6700, 0.7327, 0.5879, 0.8547, 0.8070, 0.5961, 0.7454, 0.5147, 1.2197, 0.4154, 0.4856
after torch sft-mm-after-torch-20 6.3 GiB 1.0363, 1.3195, 0.9791, 0.5142, 1.0861, 0.8911, 1.0728, 0.9029, 0.8784, 0.6701, 0.7335, 0.5883, 0.8535, 0.8060, 0.5962, 0.7453, 0.5149, 1.2208, 0.4153, 0.4853

MoE 35B 16k FA3 (EP=8, 8 GPUs)

  • MoE 35B 16k FA3 online W&B H100 SFT validation used Qwen/Qwen3.6-35B-A3B, examples/vlm_sft_moe/sft.toml, EP=8, 8 GPUs, PrimeIntellect/INTELLECT-3-SFT-10K (default / math), batch_size=8, micro_batch_size=1, seq_len=16384, pack_function="cat", shuffle=false, seed=0, max_steps=10, and attn='flash_attention_3' confirmed in resolved logs. Run root: /home/research/prime-runs/moe35b-16k-fa3-wandb-20260628T020950Z. 16k packing evidence is in packing_inspect_16k/console.log, including multi-segment packs such as [5223, 3451, 4152].
version loss_impl W&B peak mem completion / memory status losses / timings
before liger_fused before-liger n/a blocked: VLM renderer SFT unsupported before this PR n/a
before torch before-torch n/a blocked: VLM renderer SFT unsupported before this PR n/a
after liger_fused moe35b-16k-fa3-after-liger 17.7 GiB completed all 10 steps; supervisor ended CHILD_EXIT_STATUS=0, LAST_STEP=9 0.3759 (2m42s), 0.3678 (1m01s), 0.2876 (1m01s), 0.3733 (37.3s), 0.3381 (35.8s), 0.2798 (32.3s), 0.3525 (13.7s), 0.2875 (36.6s), 0.3917 (33.8s), 0.3229 (17.3s)
after torch rerun, supervised 59.1 GiB did not cleanly finalize: reached step 8 twice, then related processes exited/disappeared before step 9/final W&B sync after virtual_memory_safe_pct warnings rose to 99.2% 0.3763 (1m39s), 0.3682 (48.5s), 0.2879 (39.4s), 0.3733 (26.5s), 0.3380 (23.4s), 0.2798 (15.4s), 0.3528 (16.8s), 0.2877 (28.2s), 0.3919 (30.5s)

Logs: after_liger_seq16384_fa3/console.log and after_torch_supervised_seq16384_fa3/console.log. Final cleanup check showed all GPUs at 0 MiB and no matching SFT/torchrun processes.

Real-text E2E matrix (head 3992e110a)

  • Reran text-only SFT rows on PrimeIntellect/INTELLECT-3-SFT-10K instead of the earlier generated repeated-alpha JSONL.
  • Dataset schema validation: the HF dataset exposes splits chat, code, if, math, science, and tool with prompt / completion lists, not messages. For SFT loader compatibility, a 576-row normalized JSONL slice was generated under the run root by concatenating prompt + completion into OpenAI-style messages and dropping heterogeneous tools / tool_calls fields.
  • Current PR head: 3992e110a9faf665061bef149ddb7ef643b45ae1. Pre-PR HF baseline: merge-base 436873190d0b6d2b0a934075184d85f25952350d.
  • Multimodal 1-GPU follow-up rows stayed on the existing plex_traces_sample_256 dataset. CP is N/A for 1-GPU rows.

Current PR Real-Text Rows

row GPUs CP impl loss attn dataset status W&B
MoE text CP 8 2 custom torch flash_attention_3 PrimeIntellect/INTELLECT-3-SFT-10K PASS wmhe4mh6
Dense text CP 8 2 custom torch flash_attention_3 PrimeIntellect/INTELLECT-3-SFT-10K PASS b8vvbbuj
MoE text no-CP 8 1 custom liger_fused flash_attention_3 PrimeIntellect/INTELLECT-3-SFT-10K PASS (SFT trainer finished!; post-finish launcher teardown required cleanup) d5gltr4u
Dense text no-CP HF 8 1 hf liger_fused flash_attention_3 PrimeIntellect/INTELLECT-3-SFT-10K PASS 2yg41duy
Dense text no-CP custom 8 1 custom liger_fused flash_attention_3 PrimeIntellect/INTELLECT-3-SFT-10K PASS 44xe16l4

Dense 1-GPU Follow-Up

row GPUs impl modality dataset status W&B
Dense text no-CP HF 1 hf text PrimeIntellect/INTELLECT-3-SFT-10K PASS q7p4vjmg
Dense text no-CP custom 1 custom text PrimeIntellect/INTELLECT-3-SFT-10K PASS qkncu27w
Dense multimodal no-CP HF 1 hf multimodal plex_traces_sample_256 PASS hahunt82
Dense multimodal no-CP custom 1 custom multimodal plex_traces_sample_256 PASS tvb2yy7z

Pre-PR HF Baseline

baseline SHA GPUs impl dataset status W&B
Pre-PR dense text no-CP HF 436873190d0b6d2b0a934075184d85f25952350d 8 hf PrimeIntellect/INTELLECT-3-SFT-10K PASS tifuvcxk
Pre-PR dense text no-CP HF 436873190d0b6d2b0a934075184d85f25952350d 1 hf PrimeIntellect/INTELLECT-3-SFT-10K PASS aeyl8wms

Packing benchmarks (cat vs stack)

Benchmark results:

seq_len pack step avg actual tok/s train tok/s avg actual/train tok segments/images/pixel_rows peak mem
16k stack 0.760s 16,114 736 12,239 / 559 1.00 / 2.58 / 11,265 4.7 GiB
16k cat 0.782s 16,652 747 13,027 / 584 1.05 / 2.63 / 11,495 4.7 GiB
32k stack 0.913s 16,883 967 15,418 / 884 1.00 / 4.00 / 17,472 7.6 GiB
32k cat 1.582s 17,841 1,103 28,223 / 1,746 1.74 / 8.53 / 37,243 7.6 GiB
65k stack 1.158s 15,982 1,207 18,501 / 1,397 1.00 / 5.84 / 25,518 12.5 GiB
65k cat 3.279s 17,301 1,159 56,727 / 3,801 3.00 / 16.74 / 73,107 12.6 GiB

Fixed-work E2E results:

seq_len / setup pack steps wall step time process time samples/s actual tok/s train tok/s
32k, same 64 samples stack 64 66.231s 103.476s 0.966 17,482 1,077
32k, same 64 samples cat 47 64.623s 101.798s 0.990 17,917 1,104
65k no-warmup comparable first 64 samples stack 64 77.945s 113.210s 0.821 16,899 1,117
65k no-warmup comparable first 64 samples cat 25 77.000s 110.380s 0.831 17,106 1,130

Benchmark conclusion: cat wins by real/actual-token work and gives a modest fixed-work E2E win. Logged padded throughput can be misleading for stack because padded tokens inflate the reported token/s number.

Benchmark caveat: FLA_TILELANG was unset, and logs reported the FLA fast-path fallback because causal_conv1d was missing, so these results should be read with that environment limitation in mind.

Operational Notes

  • cat is again the default SFT packing strategy. Use stack deliberately when benchmarks for a specific dataset/model/loss combination show it helps.

Note

High Risk
Large changes to SFT tokenization defaults, multimodal forward/checkpoint paths, and distributed CP behavior; misconfiguration can hang NCCL or silently change loss masks vs legacy SFT.

Overview
SFT is renderer-only end-to-end: tokenization and loss masks always go through renderers.build_training_sample; the incremental apply_chat_template path is removed. [renderer] defaults to auto and DefaultRenderer is rejected for real data (fake data can still benchmark without a renderer). CI reverse-text configs and docs now require explicit renderers (e.g. qwen3).

Multimodal SFT adds mm_kwargs, mm_token_type_ids, and seq_lens through dataset packing, collation, and forward(). The trainer loads an AutoProcessor, wires it into the renderer, saves it with VLM checkpoints, and enforces VLM rules (custom PrimeRL VLM impl, micro_batch_size = 1, image-only, LoRA vs frozen vision). data_files supports local JSONL/zstd with normalization; CatDataset pads packs to seq_len and can mix text and image samples. New examples/vlm_sft / vlm_sft_moe configs are included.

Qwen3.5 custom models add dense Qwen3_5ForCausalLM (HF vision + PrimeRL text) alongside MoE CP fixes: set_context_parallel_attributes + setup_model_cp, global seq_lens for GatedDeltaNet/CP, and VLM ulysses sharding inside the model. CP hardening validates shard divisibility and pads cat packs so ulysses all-to-all cannot deadlock on odd lengths.

Config defaults shift (loss_implliger_fused, default seq_len 256) and validators cover Qwen3.5 stack+flash, VLM constraints, and dropping the old “renderer incompatible with VLM” rule.

Reviewed by Cursor Bugbot for commit e237511. Bugbot is set up for automated code reviews on this repo. Configure here.

@hallerite hallerite self-requested a review May 13, 2026 11:20
@eligotts eligotts force-pushed the feat/sft-multimodal-renderer branch from 3256a0d to ef07356 Compare May 22, 2026 04:12
@eligotts eligotts force-pushed the feat/sft-multimodal-renderer branch from 66a7ce6 to 4fea65d Compare May 29, 2026 06:04
@hubert-marek hubert-marek force-pushed the feat/sft-multimodal-renderer branch from b3f7123 to cc1c29a Compare May 29, 2026 06:52
@hubert-marek

Copy link
Copy Markdown
Contributor

@codex

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 3442a7372b

ℹ️ 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".

Comment thread src/prime_rl/trainer/sft/data.py Outdated
Comment on lines +665 to +668
"input_ids": torch.tensor(sample["input_ids"], dtype=torch.long, device="cuda"),
"position_ids": torch.tensor(sample["position_ids"], dtype=torch.long, device="cuda"),
"loss_mask": torch.tensor(sample["loss_mask"], dtype=torch.bool, device="cuda"),
"target_ids": torch.tensor(sample["target_ids"], dtype=torch.long, device="cuda"),

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Add a batch dimension for stacked VLM samples

When data.pack_function = "stack" and a sample has mm_kwargs, StackDataset now bypasses bucketing and yields the raw single sample, whose fields are 1-D lists. This collate path converts those directly to rank-1 tensors, unlike the packed text path's list-of-lists, so VLM SFT with stack packing sends input_ids shaped [seq] into forward()/compute_loss() where the rest of the trainer expects [batch, seq], causing shape errors or malformed model inputs.

Useful? React with 👍 / 👎.

self.logger.warning(
f"Did not find EOS token ID {self.tokenizer.eos_token_id} in input_ids. Is something wrong with the chat template? Manually appending EOS token..."
)
input_ids.append(cast(int, self.tokenizer.eos_token_id))

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Leave room for EOS after VLM truncation

When a rendered VLM example is longer than seq_len + 1 and the safe cut lands exactly at that budget, the truncated prefix usually no longer contains the original EOS, so this append makes the sequence one token too long. After the causal shift, input_ids has length seq_len + 1, and because multimodal samples bypass Cat/Stack truncation, the trainer forwards a batch that exceeds the configured context window instead of dropping or trimming it.

Useful? React with 👍 / 👎.

Comment thread src/prime_rl/trainer/sft/data.py Outdated
Comment thread src/prime_rl/trainer/sft/data.py Outdated
Comment thread src/prime_rl/trainer/sft/data.py Outdated
@hubert-marek

Copy link
Copy Markdown
Contributor

@codex

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: e2f3635fd7

ℹ️ 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".

Comment thread src/prime_rl/trainer/sft/data.py Outdated
Comment on lines +799 to +801
tools = new_record.get("tools")
if isinstance(tools, (list, dict)):
new_record["tools"] = json.dumps(tools)

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Normalize tool_defs during local JSON loading

When local data_files use the verifiers rollout shape with tool_defs instead of tools, this normalization leaves the nested tool schema as a list/dict. The SFT parser later explicitly accepts tool_defs (example.get("tools", example.get("tool_defs"))), and those per-row JSON schemas are the same Arrow-incompatible data this helper is trying to stringify. With two JSONL rows whose tool_defs[*].parameters differ, load_dataset("json", data_files=...) can still fail before _process runs because only tools is converted to a JSON string here; apply the same conversion to tool_defs.

Useful? React with 👍 / 👎.

Comment thread src/prime_rl/trainer/sft/train.py
f"{p.resolve()}:{stat.st_mtime_ns}:{stat.st_size}:{multimodal}".encode()
).hexdigest()[:12]
tmp = Path(tempfile.gettempdir()) / f"{p.stem}.{digest}.prime_rl_normalized.jsonl"
if is_writer and (not tmp.exists() or tmp.stat().st_size == 0):

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Rewrite partial temp outputs atomically

If a run is interrupted while rank-local normalization is writing this $TMPDIR file, it can leave a non-empty truncated JSONL behind. On the next run with the same source path/mtime/size digest, this condition treats that partial output as valid and skips regeneration, so HF can load corrupt or incomplete training data; write to a separate temporary path and atomically rename (or always overwrite) before considering the cache reusable.

Useful? React with 👍 / 👎.

Comment thread src/prime_rl/trainer/sft/data.py Outdated
eligotts and others added 7 commits June 25, 2026 23:22
_flatten_mm_items dropped the renderer's pixel_values/image_grid_thw
because they arrive as numpy arrays (return_tensors="np") while the
filter only kept torch.Tensor — so mm_kwargs ended up empty and VLM
SFT silently trained text-only. Coerce array-likes via torch.as_tensor.

compute_loss now passes mm_kwargs=/mm_token_type_ids= by name to the
generalized forward() (matching the RL trainer) instead of splatting
pixel_values, which the current forward() signature rejects.

Co-authored-by: Cursor <cursoragent@cursor.com>
Remove the unused `is_multimodal` attribute (and its renderer import) and
the unread `_solo` flag on bypassed multimodal samples; import
`PlaceholderRange` that was referenced only in an annotation.

Co-authored-by: Cursor <cursoragent@cursor.com>
Single-node 8-GPU MoE VLM SFT example (expert parallelism + LoRA) using
selective activation checkpointing that excludes routed_experts to avoid
the MoE router-recompute CheckpointError.

Co-authored-by: Cursor <cursoragent@cursor.com>
mm_token_type_ids flags the modality of each input token, so after the
causal shift it must mirror input_ids ([:-1]). It was sliced [1:], which
shifted every modality label one token left — marking the token before a
media run as media and the last media token as text.

Co-authored-by: Cursor <cursoragent@cursor.com>
- stack_collate: add the batch dim for solo multimodal samples so VLM SFT
  with pack_function="stack" yields [batch, seq] like the text path.
- VLM truncation budget = seq_len (was seq_len + 1) to leave room for a
  re-appended EOS after the causal shift, keeping un-packed VLM samples
  within seq_len.
- _resolve_local_data_files: disambiguate temp files by a digest of the
  source path + mtime + size (+ multimodal flag) so same-basename or
  changed/stale sources don't collide or get silently reused.
- Materialize local data per node (local-rank-0 gate) so multi-node runs
  don't read a normalized file only written on node 0's /tmp.
- Only wrap string message content into OAI block lists for multimodal
  data; text-only SFT keeps string content for the legacy chat-template path.

Co-authored-by: Cursor <cursoragent@cursor.com>
…mmetry

A text-only micro-batch (no images, or all images truncated away) made its
rank skip the vision encoder's FSDP all-gather, desyncing collectives across
ranks under FSDP + expert parallelism and deadlocking the run (NCCL timeout).
Always invoke the frozen vision encoder, feeding a tiny dummy image when none
is present, and skip the masked_scatter in that case so LM embeddings are
untouched.

Co-authored-by: Cursor <cursoragent@cursor.com>
@hubert-marek hubert-marek force-pushed the feat/sft-multimodal-renderer branch from 4138ad7 to c52ebcc Compare June 25, 2026 23:23
@hubert-marek

Copy link
Copy Markdown
Contributor

@codex

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: d431688c12

ℹ️ 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".

f"Did not find EOS token ID {self.tokenizer.eos_token_id} in input_ids. Is something wrong with the chat template? Manually appending EOS token..."
)
input_ids.append(cast(int, self.tokenizer.eos_token_id))
loss_mask.append(True)

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P1 Badge Avoid training EOS at artificial VLM truncation

When a multimodal sample exceeds seq_len, the new image-aware truncation can remove the real EOS and all assistant tokens before this block runs; this line then appends a trainable EOS, so long prompt/image-heavy examples that should be truncated or skipped instead teach the model to terminate at the context boundary. This is especially damaging for VLM traces where screenshots push the answer past seq_len; the appended EOS also makes the later sum(loss_mask[:seq_len]) check pass even if no real trainable token survived.

Useful? React with 👍 / 👎.

:12
]
tmp = Path(tempfile.gettempdir()) / f"{p.stem}.{digest}.prime_rl_normalized.jsonl"
if is_writer and (not tmp.exists() or tmp.stat().st_size == 0):

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Replace normalized files atomically

If normalization or decompression is interrupted after opening tmp for writing, the file remains non-empty, and the next run skips regeneration because this condition only rejects missing or zero-byte outputs. That can make load_dataset consume a truncated/corrupt normalized JSONL from a previous failed run; write to a unique temporary path and replace() it only after _normalize_jsonl completes.

Useful? React with 👍 / 👎.

Comment thread packages/prime-rl-configs/src/prime_rl/configs/sft.py Outdated
Comment thread src/prime_rl/trainer/sft/train.py
Comment thread src/prime_rl/trainer/models/qwen3_5_moe/mrope.py
Comment thread src/prime_rl/trainer/sft/data.py Outdated
Comment on lines +808 to +810
tools = new_record.get("tools")
if isinstance(tools, (list, dict)):
new_record["tools"] = json.dumps(tools)

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Normalize tool_defs in local JSONL

Local JSONL rows that use the supported tool_defs column still pass a list/dict with per-tool schemas through this normalizer because only tools is stringified. With verifiers-style data, PyArrow can fail schema inference before SFT starts when parameters differ across rows; apply the same JSON stringification to tool_defs as well.

Useful? React with 👍 / 👎.

Comment thread src/prime_rl/trainer/sft/data.py Outdated
@hubert-marek hubert-marek self-assigned this Jun 26, 2026
@hubert-marek hubert-marek marked this pull request as ready for review June 26, 2026 04:04
Comment thread src/prime_rl/trainer/sft/data.py

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 9c4c937bb2

ℹ️ 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".

Comment thread src/prime_rl/trainer/sft/data.py Outdated
for path in data_files:
p = Path(path)
if p.suffix in (".zst", ".jsonl"):
stat = p.stat()

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Expand data_files globs before statting them

For local JSON datasets configured with a glob such as data_files = ["data/*.jsonl"], this path has suffix .jsonl, but Path.stat() is called on the literal glob string before Hugging Face load_dataset gets a chance to expand it. Those common sharded-file configs now fail with FileNotFoundError; expand glob patterns first or leave them to load_dataset when they cannot be statted as concrete files.

Useful? React with 👍 / 👎.

Comment thread docs/training.md Outdated
@@ -1200,6 +1218,8 @@ def forward(
# via the mm_kwargs shape so we don't enumerate model_types.
if "image_grid_thw" not in mm_kwargs:

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Treat video grids as MRoPE inputs

For Qwen-style video-only samples the renderer supplies video_grid_thw/pixel_values_videos but no image_grid_thw, so this check takes the non-MRoPE path, forwards the trainer's 1D position_ids, and withholds seq_lens. Those batches then cannot build correct 3D video positions (or fail fast for unsupported video), causing video placeholder tokens to be handled like ordinary text positions; include video_grid_thw in the MRoPE detection.

Useful? React with 👍 / 👎.

pack["mm_kwargs"] = {key: value for key, value in sample_mm_kwargs.items()}
else:
for key, value in sample_mm_kwargs.items():
pack["mm_kwargs"][key] = torch.cat([pack["mm_kwargs"][key], value], dim=0)

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Merge sparse multimodal kwarg sets safely

When cat packing combines multimodal samples whose processors emit different kwarg keys, such as an image-only sample followed by a video-only sample (pixel_values/image_grid_thw vs pixel_values_videos/video_grid_thw), this indexes a key that is absent from the existing pack and raises KeyError. Initialize missing keys instead of assuming every later sample has exactly the same multimodal kwargs.

Useful? React with 👍 / 👎.

Default SFT to stack packing after VLM benchmarks showed cat packing is slower with fused loss, while preserving separated multimodal cat packing for explicit opt-in.

Co-authored-by: Cursor <cursoragent@cursor.com>
Comment thread src/prime_rl/trainer/sft/data.py Outdated
hubert-marek and others added 11 commits July 3, 2026 03:52
The generic HF VLM path has no FSDP-safe handling of mixed text/image
batches (dummy vision pass) and no pack-boundary support, so it hangs or
silently cross-attends. VLM training now requires a custom PrimeRL VLM
implementation; subsumes the previous VLM+CP-only guard.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Video-only samples bypassed the MRoPE video guard (it lives in the image
branch) and trained with 1D positions while the video tensors were
ignored. Raise at sample processing instead. Also apply ruff formatting
missed in a prior commit.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
sdpa/eager attn with the default cat packing left cu_seqlens=None for
2D (text-only) packs, since the branch only fired for 3D MRoPE
position_ids. GatedDeltaNet state and full-attention masks then
crossed packed-sample boundaries silently. Generalize the existing
seq_lens branch to both cases; multi-document packs with
full_attention layers now raise (matching the existing MRoPE
guard) instead of corrupting.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Qwen3_5GatedFlashAttention's varlen path assumes batch size 1
(cat packing); a stack bucket with more than one sample sends only
query_states[0] to the flash kernel while cu_seqlens/output_proj
still expect the full batch, crashing with a shape mismatch. Reject
the combination at config validation, matching the existing
stack+CP rejection.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
raise ValueError(
"Qwen3.5 custom flash attention does not support pack_function='stack' for validation data; use 'cat'"
)
return self

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Stack flash guard misses models

Medium Severity

validate_qwen3_5_stack_flash only treats a run as Qwen3.5 when model.name contains Qwen3.5 or qwen3_5, so configs for checkpoints whose hub name omits that string (for example Qwen/Qwen3.6-35B-A3B in examples/vlm_sft_moe) skip the pack_function='stack' rejection even when impl is custom and flash attention is enabled. Training can then hit the known flash varlen shape mismatch instead of failing at config parse time.

Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit 5fa1816. Configure here.

hubert-marek and others added 3 commits July 3, 2026 05:06
_build_cp_context decided global-vs-local cu_seqlens by reading
cu_seqlens[-1].item(), forcing a pipeline stall per linear-attention
layer per microbatch under CP. The root forward knows the provenance
statically; thread it down as a bool.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…emplate kwargs, just run-wide not per example

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Cursor Bugbot has reviewed your changes and found 2 potential issues.

There are 3 total unresolved issues (including 1 from previous review).

Fix All in Cursor

❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.

Reviewed by Cursor Bugbot for commit e237511. Configure here.

raise ValueError(
"VLM SFT requires val.data.micro_batch_size = 1 (image samples can't be packed across samples)."
)
return self

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

VLM stack packing unchecked

Medium Severity

VLM SFT validates micro_batch_size = 1 but not pack_function = "cat". With pack_function = "stack", multimodal samples bypass bucketing and are yielded at full rendered length without the CatDataset padding to seq_len, so CP sharding and fixed-length training assumptions no longer hold.

Additional Locations (1)
Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit e237511. Configure here.


@classmethod
def is_prime_state_dict(cls, state_dict: dict[str, Tensor]) -> bool:
return True

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

State dict always matches

Low Severity

Qwen3_5PreTrainedModel.is_hf_state_dict and is_prime_state_dict both always return True. Weight-load logic in get_model treats snapshots as both HF and PrimeRL compatible whenever those hooks are consulted, which can skip needed conversion or trigger redundant one-time convert paths incorrectly.

Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit e237511. Configure here.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: e2375119f3

ℹ️ 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".

Comment on lines +265 to +266
if seq_lens.numel() > 1 and "full_attention" in self.config.layer_types:
raise ValueError("Packed Qwen3.5 batches with full_attention layers require flash attention")

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Ignore synthetic padding when rejecting SDPA packs

When model.attn = "sdpa"/"eager" is used with Qwen3.5 full-attention layers and pack_function = "cat", an ordinary underfilled single-sample pack is padded by CatDataset as seq_lens=[sample_len, pad_len]. This condition treats that trailing loss-masked padding as a second packed sample and raises, so valid runs fail on common overflow/final packs even though no real token can attend across the pad boundary; filter out the synthetic padding segment before applying this rejection.

Useful? React with 👍 / 👎.

Comment on lines +831 to +832
if multimodal and isinstance(content, str):
m["content"] = [{"type": "text", "text": content}]

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Normalize prompt/completion local multimodal rows

When a local data_files JSONL uses the supported prompt/completion layout for multimodal SFT, this string-to-text-block rewrite only runs for the messages field; prompt[].content and completion[].content can still mix plain strings with image content arrays across rows. Those files can hit PyArrow schema errors in load_dataset("json", ...) before _process() can call normalize_messages, so local VLM data only works reliably in the messages layout; apply the same normalization to prompt/completion or reject that layout for multimodal local files.

Useful? React with 👍 / 👎.

Comment on lines +186 to +187
if config.data.type != "fake":
renderer = setup_renderer(tokenizer, config.renderer)

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Build a renderer when fake training has validation

When data.type = "fake" is combined with a validation config, this leaves renderer as None, but val.data is always real SFT data and run_validation() later calls setup_dataset(..., renderer=renderer), which raises SFT data requires a renderer before the first eval. Construct the renderer whenever either training data is real SFT or validation is configured.

Useful? React with 👍 / 👎.


logger.info(f"Initializing tokenizer ({config.tokenizer})")
tokenizer = setup_tokenizer(config.tokenizer)
processor = setup_processor(config.tokenizer)

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Load VLM processors from the model checkpoint

When [tokenizer] name points at a tokenizer-only path while [model] name is the VLM checkpoint, this probes AutoProcessor at the tokenizer path; setup_processor() then returns None and the [model.vlm] check aborts even though the model checkpoint has the image processor needed for rendering. Load the processor from model.name or add an explicit processor config instead of tying it to the tokenizer override.

Useful? React with 👍 / 👎.

gen_config.save_pretrained(path)
tokenizer.save_pretrained(path)
if processor is not None:
processor.save_pretrained(path)

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Save the updated tokenizer after the processor

When a VLM run uses a custom [tokenizer] chat_template or relies on the tokenizer mutation in setup_tokenizer, saving the separately loaded processor after tokenizer.save_pretrained() can overwrite the exported tokenizer files and chat template because ProcessorMixin.save_pretrained() also saves its embedded tokenizer. Save the processor first or attach the updated tokenizer to the processor before saving so weight checkpoints preserve the tokenizer actually used for training.

Useful? React with 👍 / 👎.

@hubert-marek

Copy link
Copy Markdown
Contributor

Closing in favor of a 7-PR stack that splits this work into reviewable slices (per @samsja's suggestion), with #2889's content folded in:

  1. feat(sft): renderer-only SFT tokenization + local data files #2942 — feat(sft): renderer-only SFT tokenization + local data files
  2. feat(models): dense Qwen3.5 custom attention #2943 — feat(models): dense Qwen3.5 custom attention
  3. feat(trainer): seq_lens as the packed-sample boundary contract (SFT + RL) #2944 — feat(trainer): seq_lens as the packed-sample boundary contract (SFT + RL) (includes feat(trainer): pack RL multimodal samples #2889's seq_lens contract)
  4. feat(sft): multimodal (VLM) SFT #2945 — feat(sft): multimodal (VLM) SFT
  5. feat(cp): context parallelism for hybrid + VLM SFT #2946 — feat(cp): context parallelism for hybrid + VLM SFT (this PR's tip is tree-equivalent to this branch, modulo review fixes enumerated in the PR bodies)
  6. feat(rl): pack multimodal samples with context parallelism #2947 — feat(rl): pack multimodal samples (feat(trainer): pack RL multimodal samples #2889's packing mechanics, always-on)
  7. feat(rl): multimodal context parallelism #2948 — feat(rl): multimodal context parallelism (new code — RL MM CP existed in neither source PR)

Each PR is based on the previous one and states its dependencies; all review feedback from this PR landed in the corresponding slice. The stack also adopts the custom-only VLM training policy (hf-VLM training is dropped; the MM configs migrate to Qwen3.5).

samsja added a commit that referenced this pull request Jul 9, 2026
SFT tokenization now goes exclusively through the renderers library:
- build_training_sample with a hand-coded renderer replaces the
  incremental apply_chat_template masking path, which corrupted loss
  masks under position-dependent chat templates (e.g. Qwen3 stripping
  past thinking blocks across user turns). DefaultRenderer is rejected
  in setup_renderer; fake data skips renderer construction entirely.
- [renderer] config defaults to auto-resolution from the tokenizer;
  template controls (e.g. enable_thinking) are set run-wide via the
  typed renderer config.
- Non-assistant roles opt into the loss via the renderer's body-only
  path (content_sft_roles).
- _drop_null_fields strips Arrow phantom nulls before tool-call
  deserialization so genuine nulls inside argument strings survive.
- Null-check rather than key-check for messages vs prompt/completion
  resolution (Arrow schema union adds messages: null to prompt/
  completion rows).
- seq_len default 256 (FakeDataConfig keeps 128); loss_impl defaults
  to liger_fused; CI reverse-text configs pin [renderer] name=qwen3.

Split from #2942 (part 1/7 of #2485), renderer-only without data_files.
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.

3 participants