feat(trainer): seq_lens as the packed-sample boundary contract (SFT + RL)

[hubert-marek]

Summary

Part 3/7 of the split of #2485 + #2889. Based on #2943 (feat/qwen35-dense-attention) — needs the dense Qwen3.5 model for its seq_lens consumption; the renderer-only SFT data path comes from #2942.

seq_lens becomes the packed-sample boundary contract for both trainers; models consume it for varlen attention.

SFT + models (first commit):

• CatDataset rework: _finalize_pack pads every pack to seq_len (pad tokens form their own loss-masked document, so packs are always full-length and CP-divisible), yields the tail pack instead of dropping it, and overflow samples start the next pack instead of being truncated mid-sample.
• Sample/Batch gain seq_lens; cat_collate forwards it for batch-size-1 packs; stack bucketing yields seq_lens=None.
• seq_lens is a universal contract: every custom ForCausalLM declares the typed seq_lens forward parameter and the trainer passes it unconditionally to any PreTrainedModelPrimeRL (generic HF models never receive it). Standard-attention families keep deriving cu_seqlens from position_ids on the flash path (boundary-equivalent — the packer resets positions per document) and now reject packed rows loudly on non-flash paths instead of silently attending across documents; gpt_oss/glm_moe_dsa honor boundaries via position_ids (HF packed-sequence masks / sparse MLA varlen indices) and accept the param; Qwen3.5 dense + MoE consume it for varlen attention and fla state resets.
• Dense + MoE Qwen3.5 build cu_seqlens from seq_lens on the non-flash path. Previously cu_seqlens stayed None there: fla's GatedDeltaNet carried conv/SSM state across packed samples and SDPA attended across document boundaries — silent cross-sample contamination for sdpa/eager + cat packing. Packed hybrid batches with full_attention layers require flash and raise otherwise.
• Config validator rejects pack_function="stack" for Qwen3.5 custom flash attention (its varlen kernel assumes batch 1; stack produces genuine batches).
• Under CP both trainers pass seq_lens=None for now — boundaries span the pre-shard pack; global semantics land in feat(cp): context parallelism for hybrid + VLM SFT #2946.
• The HF-impl monkey-patch keeps deriving cu_seqlens from position_ids only; it never receives seq_lens (only custom PreTrainedModelPrimeRL models do), so no seq_lens branch is added there.

RL (second commit, from #2889): packed RL micro-batches carry the same boundaries end to end — prepare_sample stamps them, _materialize_bin accumulates and asserts them, pad_micro_batch records padding as its own segment, the RL dataloader tensorizes them into forward().

Data efficiency: old vs new packing (offline replay, 12k docs, real tokenizer path, seq 8192)

The dataset (long CoT, mean doc 6,783 tokens) makes the trade stark:

	old packing (main)	new packing (this PR)
tokens trained per epoch	41.4M	52.9M (+27.7%)
data discarded per epoch	49.1% (chopped remainders)	0% (only the inherent overlong cap)
docs ending mid-sample	42.1%	0%
pad overhead per row	0%	21.4%

The old packer's "zero padding" chops nearly every pack's last document and discards its remainder — half the dataset's tokens are never trained, and 42% of docs train without their endings/EOS. The new packer pays 21.4% pad compute (hardware tok/s is unchanged: pads flow through as a masked document) to train 28% more tokens with correct boundaries. Pad fraction shrinks on shorter-doc datasets.

E2E (W&B project pr2485)

• Text SFT dense + MoE with reworked cat packing: dense 9gjfrcx5, MoE jyaqjzo1 — 10/10 steps each
• RL text run at this tip (seq_lens end-to-end through transport → packer → dataloader → forward): 10/10 steps — bec3d2415a
• Integration suite (SFT+RL subset) at this tip: 7/8 passed; the one failure is test_loss_goes_down — a marginal 5-step assertion (loss delta +0.0018) whose values match the feat(sft): Enable VLM SFT  #2485 reference bit-for-bit; see feat(rl): multimodal context parallelism #2948 for the full-suite/main-control analysis

🤖 Generated with Claude Code

---

Note

High Risk
Touches core training forward paths, packing semantics, and many custom model implementations; incorrect boundaries would silently corrupt attention or loss, and packed non-flash configs now fail loudly instead of cross-attending.

Overview
Introduces seq_lens as the shared packed-document boundary metadata from dataloaders through forward(), so attention and hybrid layers can respect document boundaries instead of leaking across packed samples.

SFT packing reworks CatDataset: overflow starts a new pack instead of truncating mid-sample, tail packs are yielded, and each row is padded to fixed seq_len with pad tokens as their own masked segment (with seq_lens tracking real vs pad lengths). RL mirrors the same contract in prepare_sample, bin materialization, padding, and the RL batch type.

The trainer passes seq_lens only to custom PreTrainedModelPrimeRL models; generic HF models never receive it. Most custom stacks still build varlen indices from position_ids on flash but error on multi-segment packs without flash; Qwen3.5 dense/MoE can build cu_seqlens from seq_lens on the non-flash path. Context parallelism clears seq_lens in SFT/RL for now. CI reverse-text SFT configs bump step counts to match longer runs.

^{Reviewed by Cursor Bugbot for commit ededf23. Bugbot is set up for automated code reviews on this repo. Configure here.}

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: d1e5538558

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

[chatgpt-codex-connector]

Persist overflow sample before yielding

When seq_len + sample_len overflows, this yield happens after the next sample has already been pulled from the wrapped dataset but before it is appended to the next pack. If StatefulDataLoader checkpoints immediately after returning that batch, state_dict() only records the wrapped dataset state, so the fetched overflow sample is considered consumed but is not stored anywhere; resuming from that checkpoint skips one sample per worker at these overflow boundaries and changes the training data stream.

Useful? React with 👍 / 👎.

[hubert-marek]

Doesn't matter

[chatgpt-codex-connector]

Avoid treating padding as a packed sample

Appending pad_len here makes a single real short example look like multiple packed segments, e.g. seq_lens=[100, 28] for a 128-token pack with trailing padding. Custom Qwen3.5 non-flash models reject any seq_lens.numel() > 1 when full_attention layers are present, so otherwise valid single-sample padded packs such as final validation batches or overflow flushes fail unless users switch to flash attention, even though the extra segment is only loss-masked padding.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 8740508afa

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

[chatgpt-codex-connector]

Avoid adding padding as a Qwen segment

When RL padding is enabled (pad_to_multiple_of > 1) and a Qwen3.5 custom non-flash model has any full_attention layers, a single otherwise valid sample whose length is not already aligned gets seq_lens like [real_len, padding_size]. The Qwen3.5 forward path treats seq_lens.numel() > 1 as a packed batch and raises Packed Qwen3.5 batches with full_attention layers require flash attention, so this turns alignment padding into a hard runtime failure even though no second real sample was packed.

Useful? React with 👍 / 👎.

[mikasenghaas]

why this change?

[hubert-marek]

We don't want to truncate mm_kwargs, and mm_token_type_ids, and seq_lens - only the ("input_ids", "position_ids", "loss_mask", "target_ids")

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

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

^{Reviewed by Cursor Bugbot for commit ededf23. Configure here.}

[cursor]

Pad segments false-trigger packed guards

Medium Severity

_finalize_pack and pad_micro_batch always append padding as its own seq_lens segment, so almost every underfull pack has seq_lens.numel() > 1. The new non-flash guards treat that as multi-document packing and raise, even when the pack is a single real sample plus trailing pad. Under causal attention that trailing pad cannot affect real-token outputs, so valid sdpa/eager cat-packed runs fail.

Additional Locations (2)

• src/prime_rl/trainer/models/qwen3/modeling_qwen3.py#L173-L177
• src/prime_rl/trainer/batch.py#L471-L473

 

^{Reviewed by Cursor Bugbot for commit ededf23. Configure here.}