feat(sft): Enable VLM SFT #2485
Conversation
3256a0d to
ef07356
Compare
66a7ce6 to
4fea65d
Compare
b3f7123 to
cc1c29a
Compare
There was a problem hiding this comment.
💡 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".
| "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"), |
There was a problem hiding this comment.
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)) |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
💡 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".
| tools = new_record.get("tools") | ||
| if isinstance(tools, (list, dict)): | ||
| new_record["tools"] = json.dumps(tools) |
There was a problem hiding this comment.
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 👍 / 👎.
| 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): |
There was a problem hiding this comment.
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 👍 / 👎.
_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>
4138ad7 to
c52ebcc
Compare
There was a problem hiding this comment.
💡 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) |
There was a problem hiding this comment.
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): |
There was a problem hiding this comment.
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 👍 / 👎.
| tools = new_record.get("tools") | ||
| if isinstance(tools, (list, dict)): | ||
| new_record["tools"] = json.dumps(tools) |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
💡 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".
| for path in data_files: | ||
| p = Path(path) | ||
| if p.suffix in (".zst", ".jsonl"): | ||
| stat = p.stat() |
There was a problem hiding this comment.
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 👍 / 👎.
| @@ -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: | |||
There was a problem hiding this comment.
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) |
There was a problem hiding this comment.
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>
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 |
There was a problem hiding this comment.
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.
Reviewed by Cursor Bugbot for commit 5fa1816. Configure here.
_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>
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 2 potential issues.
There are 3 total unresolved issues (including 1 from previous review).
❌ 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 |
There was a problem hiding this comment.
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)
Reviewed by Cursor Bugbot for commit e237511. Configure here.
|
|
||
| @classmethod | ||
| def is_prime_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: | ||
| return True |
There was a problem hiding this comment.
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.
Reviewed by Cursor Bugbot for commit e237511. Configure here.
There was a problem hiding this comment.
💡 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".
| 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") |
There was a problem hiding this comment.
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 👍 / 👎.
| if multimodal and isinstance(content, str): | ||
| m["content"] = [{"type": "text", "text": content}] |
There was a problem hiding this comment.
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 👍 / 👎.
| if config.data.type != "fake": | ||
| renderer = setup_renderer(tokenizer, config.renderer) |
There was a problem hiding this comment.
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) |
There was a problem hiding this comment.
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) |
There was a problem hiding this comment.
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 👍 / 👎.
|
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:
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). |
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.


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
renderers.build_training_sample; the legacy incremental tokenization path is removed.[renderer]defaults toauto(hand-coded renderer required;DefaultRendereris rejected).deps/renderersadvances to99bedaf, which surfaces the multimodal payload throughbuild_training_sample.Multimodal SFT plumbing
mm_kwargs,mm_token_type_ids, andseq_lensflow through datasets, packing, collation, and trainerforward().CatDatasetpacking path serves text-only and multimodal samples: text-only spans can share a multimodal pack with zeromm_token_type_ids; compatible multimodal kwargs concatenate along the leading item dimension.pack_function = "cat"is the SFT default again (stackremains opt-in).AutoProcessor, attaches it to the renderer, saves it with VLM weight checkpoints (soAutoProcessor.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). Newexamples/vlm_sftandexamples/vlm_sft_moeconfigs (data paths are placeholders).data_filescannot be combined withsubsets/splits(config-rejected: subsets were silently ignored and non-train splits failed at load time)..jsonl/.zstdata_filesare normalized into a$TMPDIRcache beforeload_dataset; the cache (and zstd decompression output) is written to a.partsidecar 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), whichstack's genuine multi-row batches violate, crashing with a shape mismatch.sdpa/eager) attention now buildscu_seqlensfor text-onlycatpacks, not just 3D MRoPE packs: linear-attention (GatedDeltaNet) state correctly resets at pack boundaries, and multi-document packs withfull_attentionlayers 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.data_files(e.g.data/*.jsonl) are expanded before normalization instead of beingstat()-ed literally.tool_defs(verifiers rollout format) is JSON-stringified before the Arrow load, same astools— heterogeneous per-row tool schemas otherwise crash schema inference or gain phantom null fields.nullvalues inside tool-call argument JSON survive into the rendered training text.data.loss_mask.user/system/tool = trueworks 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
Qwen3_5ForCausalLMcustom implementation (HF vision tower + PrimeRL text stack), soimpl = "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.set_context_parallel_attributes, real global document boundaries forwarded viaseq_lens(seq_lens_are_global), CP-aware conv1d, and cross-rank recurrent-state exchange.setup_model_cpreplacessetup_hybrid_cp+setup_nemotron_h_cpin 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.cu_seqlensdecision is threaded down as a provenance bool instead of reading the CUDA tensor back per linear-attention layer per microbatch.CP correctness fixes (root-caused during E2E validation on this branch)
seq_lenwith odd lengths.shard_for_cpusedtorch.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_allinput sizes[2, 699, 8, 256]vs[2, 698, 8, 256]on the stuck pair). Fixes:CatDatasetpads every pack to exactlyseq_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_cpvalidates divisibility and batch shape up front, so an uneven pack now fails loudly instead of hanging the job.E2E validation at PR head
71c82e4e0Setup: 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 ofPrimeIntellect/INTELLECT-3-SFT-10K, W&B projectqwen35-custom-attention-e2e. CP rows use CP 2 (ulysses). Resolved settings inherited from currentmaindefaults:compileon, activation offloading on,optim_cpu_offload = true.Qwen/Qwen3.5-0.8Bliger_fused0.9662, 0.8727, 0.8720, 0.6853, 0.8664, 0.9094, 0.8324, 0.9404, 0.9453, 0.7762pl1ncv8tQwen/Qwen3.5-35B-A3Bliger_fused0.6106, 0.5296, 0.5446, 0.4061, 0.5366, 0.5684, 0.5607, 0.5626, 0.6025, 0.4784fa9dtakvQwen/Qwen3.5-0.8Bliger_fused1.0451, 0.7927, 0.8430, 0.7929, 0.8006, 0.9178, 0.8406, 0.9146, 0.9459, 0.7773kcqhjnbuQwen/Qwen3.5-0.8Btorch1.0451, 0.7925, 0.8432, 0.7928, 0.8009, 0.9184, 0.8406, 0.9144, 0.9462, 0.7775lc0g7nrcQwen/Qwen3.5-35B-A3Bliger_fused0.6568, 0.4835, 0.5117, 0.4930, 0.4874, 0.5857, 0.4913, 0.5689, 0.5646, 0.4920swn527chQwen/Qwen3.5-35B-A3Btorch0.6570, 0.4842, 0.5115, 0.4926, 0.4875, 0.5851, 0.4913, 0.5686, 0.5649, 0.491916czkxsrNotes:
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.torchvsliger_fusedparity: dense losses match step-for-step to ≤2e-4, MoE to ≤1e-3 (grouped-GEMM nondeterminism).torchcosts much more memory (full logits) and step time.nan_count = 0.cfd518c93,6a01021f4, and71c82e4e0.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_downfailure analysisThe failing check is a flakiness issue exposed by the renderer change, not a training regression:
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. Atlr = 1e-6withmax_norm = 1.0over 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.<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.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)
prime-job-ssh-node-5f2d5979-0, projectpacking, withQwen/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, andmax_steps=20. Beforebf3d381used renderers1933293; afterf9185a7used renderersb4c4051. Before/after losses match exactly at logged precision for both loss implementations;torchis the repo's non-Liger CE path.liger_fusedsft-mm-before-liger-201.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.4856torchsft-mm-before-torch-201.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.4853liger_fusedsft-mm-after-liger-201.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.4856torchsft-mm-after-torch-201.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.4853MoE 35B 16k FA3 (EP=8, 8 GPUs)
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, andattn='flash_attention_3'confirmed in resolved logs. Run root:/home/research/prime-runs/moe35b-16k-fa3-wandb-20260628T020950Z. 16k packing evidence is inpacking_inspect_16k/console.log, including multi-segment packs such as[5223, 3451, 4152].liger_fusedbefore-ligertorchbefore-torchliger_fusedmoe35b-16k-fa3-after-ligerCHILD_EXIT_STATUS=0,LAST_STEP=90.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)torchrerun,supervisedvirtual_memory_safe_pctwarnings rose to99.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.logandafter_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)PrimeIntellect/INTELLECT-3-SFT-10Kinstead of the earlier generated repeated-alpha JSONL.chat,code,if,math,science, andtoolwithprompt/completionlists, notmessages. For SFT loader compatibility, a 576-row normalized JSONL slice was generated under the run root by concatenatingprompt + completioninto OpenAI-stylemessagesand dropping heterogeneoustools/tool_callsfields.3992e110a9faf665061bef149ddb7ef643b45ae1. Pre-PR HF baseline: merge-base436873190d0b6d2b0a934075184d85f25952350d.plex_traces_sample_256dataset. CP is N/A for 1-GPU rows.Current PR Real-Text Rows
customtorchflash_attention_3PrimeIntellect/INTELLECT-3-SFT-10Kwmhe4mh6customtorchflash_attention_3PrimeIntellect/INTELLECT-3-SFT-10Kb8vvbbujcustomliger_fusedflash_attention_3PrimeIntellect/INTELLECT-3-SFT-10KSFT trainer finished!; post-finish launcher teardown required cleanup)d5gltr4uhfliger_fusedflash_attention_3PrimeIntellect/INTELLECT-3-SFT-10K2yg41duycustomliger_fusedflash_attention_3PrimeIntellect/INTELLECT-3-SFT-10K44xe16l4Dense 1-GPU Follow-Up
hfPrimeIntellect/INTELLECT-3-SFT-10Kq7p4vjmgcustomPrimeIntellect/INTELLECT-3-SFT-10Kqkncu27whfplex_traces_sample_256hahunt82customplex_traces_sample_256tvb2yy7zPre-PR HF Baseline
436873190d0b6d2b0a934075184d85f25952350dhfPrimeIntellect/INTELLECT-3-SFT-10Ktifuvcxk436873190d0b6d2b0a934075184d85f25952350dhfPrimeIntellect/INTELLECT-3-SFT-10Kaeyl8wmsPacking benchmarks (cat vs stack)
Benchmark results:
Fixed-work E2E results:
Benchmark conclusion:
catwins by real/actual-token work and gives a modest fixed-work E2E win. Logged padded throughput can be misleading forstackbecause padded tokens inflate the reported token/s number.Benchmark caveat:
FLA_TILELANGwas unset, and logs reported the FLA fast-path fallback becausecausal_conv1dwas missing, so these results should be read with that environment limitation in mind.Operational Notes
catis again the default SFT packing strategy. Usestackdeliberately 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 incrementalapply_chat_templatepath is removed.[renderer]defaults toautoandDefaultRendereris 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, andseq_lensthrough dataset packing, collation, andforward(). The trainer loads anAutoProcessor, 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_filessupports local JSONL/zstd with normalization;CatDatasetpads packs toseq_lenand can mix text and image samples. Newexamples/vlm_sft/vlm_sft_moeconfigs 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, globalseq_lensfor 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_impl→liger_fused, defaultseq_len256) and validators cover Qwen3.5stack+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.