refactor(models): make seq_lens a universal custom-model contract#2971
Closed
hubert-marek wants to merge 11 commits into
Closed
refactor(models): make seq_lens a universal custom-model contract#2971hubert-marek wants to merge 11 commits into
hubert-marek wants to merge 11 commits into
Conversation
CatDataset packs now carry explicit per-sample boundaries instead of relying on position_ids resets: - _finalize_pack pads every pack to seq_len (pad tokens form their own loss-masked document) and yields the tail pack instead of dropping it; overflow samples start the next pack rather than being truncated. - Sample/Batch gain seq_lens; cat_collate forwards it for batch-size-1 packs, stack bucketing yields seq_lens=None. - trainer forward() gains seq_lens and hands it to custom models via a new PreTrainedModelPrimeRL.prime_forward_kwargs hook, so PrimeRL-only kwargs never leak into generic HF model signatures. - Dense + MoE Qwen3.5 build cu_seqlens from seq_lens on the non-flash path (previously None: fla's GatedDeltaNet carried conv/SSM state across packed samples and SDPA saw one causal document). Packed hybrid batches with full_attention layers require flash and raise otherwise; a config validator rejects pack_function='stack' for Qwen3.5 custom flash attention (its varlen kernel assumes batch 1). - Under CP the trainer passes seq_lens=None for now: boundaries span the pre-shard pack and models only have local semantics. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Mirror of the SFT-side contract: packed RL micro-batches carry explicit per-sample boundaries end to end — prepare_sample stamps them, _materialize_bin accumulates them across the bin (asserting they sum to the packed length), pad_micro_batch records padding as its own segment, and the RL dataloader tensorizes them into the micro-batch for forward(). Under CP the trainer passes seq_lens=None for now, matching the SFT side, until global boundary semantics land. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Renderer-produced multimodal samples flow end to end through SFT: - SFTDataset keeps the renderer's MultiModalData: image-safe truncation (never cuts inside a placeholder run, drops truncated-out images, and skips samples whose EOS was truncated away), mm_token_type_ids kept aligned through EOS append and causal shift, video inputs rejected explicitly. _flatten_mm_items folds per-image processor outputs into model-forward kwargs (pixel_values, image_grid_thw, ...). - Cat packing concatenates mm_kwargs across samples in a pack and represents text-only spans as zeros in mm_token_type_ids; stack bucketing emits multimodal samples solo. Collates move mm tensors to CUDA and keep packed rows batch-size 1. - Local data_files normalization wraps string message content into typed content blocks under a multimodal flag so Arrow can unify text and image rows. - setup_processor loads the AutoProcessor and attaches it to the renderer; VLM SFT fails fast without one; weight checkpoints save the processor alongside the tokenizer. - Dense Qwen3.5 gains the composite VLM body (HF vision encoder + custom text model) with MRoPE positions from mm_token_type_ids and registers in _CUSTOM_VLM_MAPPING; VLM training requires a custom PrimeRL implementation. Both dense and MoE VLM bodies always run the vision encoder (dummy input on text-only microbatches, kept in the backward graph) for FSDP/EP collective symmetry. - Config: VLM SFT requires micro_batch_size=1; unfreezing the vision encoder is incompatible with LoRA; the renderer-vs-VLM rejection is gone (renderers own multimodal tokenization now). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
processor.save_pretrained rewrites tokenizer files from the processor's separately loaded tokenizer, clobbering the configured pad token and any custom chat template. Save the tokenizer last so it wins. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Without [model.vlm] the custom VLM class still loads and runs the vision tower on dummy input every microbatch, so its params carry zero — not None — grads and AdamW weight decay silently shrinks them. Freeze the tower when no VLM training is configured; mixed image/text batches cannot occur without [model.vlm], so collective symmetry is unaffected. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…l.vlm] The processor auto-loads for VLM checkpoints and mm samples would flow through packing and forward unvalidated (no mbs/attn/CP checks, no freeze policy). Require the explicit [model.vlm] declaration instead. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Unify model CP setup and make packed-boundary metadata CP-aware: - setup_model_cp replaces setup_hybrid_cp/setup_nemotron_h_cp: models owning linear-attention/Mamba layers expose set_context_parallel_attributes (dense + MoE Qwen3.5, NemotronH) and wire their own layers; hybrid models without the hook are rejected instead of silently misconfigured. - seq_lens stays global under CP: documents can straddle the shard cut, so input_ids/position_ids shard per rank while seq_lens keeps the full pre-shard boundaries, flagged via seq_lens_are_global through forward() / prime_forward_kwargs into both Qwen3.5 models. A cu_seqlens_are_global provenance flag rides along to fla's GatedDeltaNet instead of reading cu_seqlens back (which cost one CPU-GPU sync per linear-attention layer per microbatch); the CP conv path switches to fla's causal_conv1d with a CP context so conv state passes between ranks with document-boundary resets. - VLM CP defers sharding to the model: the vision encoder and the image-embed merge see the full sequence, then the root forward shards embeds/positions (shard_position_ids_for_cp handles 3D MRoPE) and updates ulysses attention params from global seq_lens. The SFT trainer skips pre-sharding for MRoPE batches; the vision encoder is pinned to SDPA under CP. - shard_for_cp validates divisibility (uneven shards deadlock ulysses all-to-all) and CatDataset's seq_len padding guarantees it; the LoRA token count follows the sharded loss_mask. - Qwen3.5 CP requires the custom implementation (the HF impl has no set_context_parallel_attributes); rejected in get_model. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Multimodal RL samples no longer force one-sample micro-batches: eager image samples pack with text spans and compatible mm_kwargs samples from the same run/LoRA, with per-sample boundaries kept in seq_lens. Packing is always on — VLM training is custom-implementation-only and the custom Qwen3.5 models advertise supports_packed_multimodal_training; a config validator requires varlen flash attention for VLM training and validate_multi_modal_pack fails loudly at startup for models without packed-MM support. Ported from #2889 minus its gating (pack_multimodal knob, resolve_pack_multimodal) and its tokenizer packer params (removed on main since). (cherry picked from commit 027edc2)
…[model.vlm] Mirrors the SFT-side check: transported mm samples would otherwise flow into forward unvalidated (no packability, attn, or freeze policy checks). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Mirror the SFT VLM CP path in the RL trainer, replacing the blanket NotImplementedError: - MRoPE multimodal batches under ulysses defer sharding to the model: input_ids/position_ids stay global so the vision encoder and image-embed merge see the full sequence; the model root forward shards embeds/positions after merge (and routed_experts with them). Text batches keep the existing pre-shard path. - seq_lens passes with seq_lens_are_global=True under CP — micro-batch boundaries span the pre-shard sequence and pad_to_multiple_of=cp already guarantees shard divisibility (padding is its own seq_lens segment). - The LoRA token-count CP adjustment derives chunk_size from the sharded length even when input_ids stays global. - Config gate relaxed: VLM + CP now requires cp_style='ulysses' and a custom model implementation instead of being rejected outright; resolve_pack_multimodal drops its CP rejection (packed MM rows meet the same divisibility guarantee). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Every custom ForCausalLM now declares the typed seq_lens / seq_lens_are_global forward parameters and the trainer passes them unconditionally; the prime_forward_kwargs hook is gone. - Standard attention families (llama, qwen3, qwen3_moe, glm4_moe, laguna, afmoe, minimax_m2, nemotron_h): flash path keeps deriving cu_seqlens from position_ids; non-flash paths now reject packed rows loudly instead of silently attending across document boundaries. - gpt_oss / glm_moe_dsa honor boundaries via position_ids (HF packed-sequence masks / sparse MLA varlen indices); they accept the params to satisfy the contract. - Qwen3.5 dense + MoE consume seq_lens as before, minus the hook. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
hubert-marek
force-pushed
the
feat/rl-multimodal-cp
branch
from
dd21849 to
3e99cac
Compare
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
1 participant
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.
Summary
Part 8 of the #2485 split stack. Based on #2948 (
feat/rl-multimodal-cp) — needs theseq_lensboundary contract from #2944 and the CP-global semantics from #2946.Makes
seq_lensa universal custom-model contract: every customForCausalLMdeclares the typedseq_lens/seq_lens_are_globalforward parameters, the trainer passes them unconditionally to anyPreTrainedModelPrimeRL, and theprime_forward_kwargshook is deleted. Previously only Qwen3.5 dense/MoE received the packed boundaries (via the hook); every other custom family silently never saw them.Per-family behavior:
cu_seqlensfromposition_ids(boundary-equivalent — the packer resets positions per document), so numerics are unchanged. Non-flash paths now reject packed rows loudly (Packed <Family> batches require flash attention) instead of silently attending across document boundaries — SDPA/eager have no varlen support.position_idson every path (HF packed-sequence masks / sparse MLA varlen indices); they accept the params to satisfy the contract, documented in the signature.seq_lensexactly as before (varlencu_seqlens, fla state resets, CP-global handling), minus the hook indirection.New-model guidance in
docs/development.mdupdated: declaring and honoring the two parameters is now part of the custom-model checklist.Validation
tests/unit/train/failure set byte-identical to the base tip's pre-existing environment baseline (59, CUDA-dependent) — no new failures; forward-contract tests updated and passing.seq_lens=[8]single-doc under sdpa and raisesPacked Llama batches require flash attentionforseq_lens=[5, 3].@auto_docstring-decorated forwards.🤖 Generated with Claude Code