* GRPO: optional sequence packing for the no-grad old/ref logp path
Add an opt-in sequence-packing fast path to _get_per_token_logps_and_entropies, enabled with
UNSLOTH_GRPO_SEQ_PACKING=1. When the batch is text-only, the padded [B, Lmax] per-chunk forward is
replaced by a single varlen [1, sum L] forward (BlockDiagonalCausalMask via packed_seq_lengths with
reset position_ids). Per-token logps use the same float32 chunked_hidden_states_selective_log_softmax
as the padded path, so the old and reference logps are bit-for-bit identical.
Safety: the packed path is self-verified once against the padded ground truth on a batch that has at
least two rows with real completion tokens (self._unsloth_seq_packing_nograd_ok), so cross-sample
contamination would actually manifest; a degenerate all-pad / fully tool-masked batch leaves the
verdict unset and re-verifies later. If a backend silently ignores packed_seq_lengths (flat batch run
under a normal causal mask, samples leaking across boundaries), the packed logps will not match and
packing is disabled instead of corrupting logps. It also forces use_cache=False (a populated
past_key_value disables varlen packing), skips packing when a sliding window is shorter than the
packed stream, runs the same GPT-OSS offload device_synchronize the padded loop uses, and falls back
on any exception (UNSLOTH_GRPO_SEQ_PACKING_DEBUG=1 prints the reason).
Default off, so existing behavior is unchanged. Pairs with the matching gradient-path change in
unsloth_zoo so the full GRPO logp + loss + backward can run packed.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* GRPO no-grad packing: address review feedback
- Cache the packed-vs-padded verdict per unwrapped model instead of on the
trainer, so a separately forwarded reference model is verified on its own
forward path rather than inheriting the policy model's verdict.
- Force the padded path when token_type_ids or mm_token_type_ids are present,
matching the extra vision kwargs the padded loop forwards.
- Require the xformers varlen backend before packing. Without it the packed
mask falls back to a dense O(T^2) SDPA mask that can OOM on the flattened
batch, so we keep the padded loop in that case.
- On any packed-forward failure (missing backend, OOM, unsupported forward)
empty the cache on OOM, disable packing for that model, and fall back to the
chunked padded loop instead of retrying every step.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* GRPO no-grad packing: default-on, verify against per-row reference
Redesign of the optional sequence-packing fast path for the no-grad
old/ref logprob recompute, after establishing that the packed forward is
the exact per-row computation and the padded batch forward is the side
that mis-positions left-padded rows on long completions.
- Default the packing on (UNSLOTH_GRPO_SEQ_PACKING, disable with 0).
- Verify the packed logprobs against the per-row clean forward (each
row's real tokens alone, reset 0-based positions, no padding), not the
padded batch which is itself wrong for left-padding. Cross-sample
contamination (a backend ignoring packed_seq_lengths) shows up as a
large mismatch and falls back to the padded loop.
- Make the trust decision shape and RoPE aware: re-verify whenever the
packed total length or the longest segment grows past what was
verified, so a later batch crossing a LongRoPE short/long cache
boundary is re-checked instead of trusted blindly.
- Run lm_head only on completion-prediction positions instead of every
packed prompt token, so long-prompt/short-completion batches do not
pay for projecting the whole packed prompt.
- Drop the hard xformers import so the path also runs in
FlashAttention-only environments; the per-row verification guards
correctness regardless of backend.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* GRPO no-grad packing: disable entirely on cross-sample mismatch
When the per-row verification fails, distinguish the two failure modes by
magnitude instead of by sequence length:
- A large mismatch (>= 1.5) is the cross-sample contamination signature:
the model's attention does not honor the block-diagonal packed mask
(seen on some MoE / custom-attention models, e.g. qwen2_moe). Disable
packing entirely for the model so later batches do not pay the
verification cost again.
- A moderate mismatch is more likely a length-boundary effect (a LongRoPE
short/long cache switch): keep marking just that length region unsafe so
packing still runs for smaller shapes.
Validated: Qwen1.5-MoE falls back after a single verification (grad and
no-grad ok flags go False, no re-verify on later steps); dense Llama-3.2
and Qwen3 still verify and engage packing.
* GRPO no-grad packing: trim comments to be concise
* GRPO no-grad packing: fix per-row completion boundary for left-padded rows
The completion-target selection used a single global boundary
(col >= L - logits_to_keep). After left-packing, each row's completion
starts at (L - logits_to_keep) - left_pad[row], so for left-padded rows
the first left_pad completion tokens fall below the global boundary and
were dropped, leaving 0 logprobs at real completion positions that the
loss mask keeps. Use the per-row boundary so packed coverage matches
create_completion_attention_mask exactly, and widen the self-verify mask
to the full per-row completion region so it can catch coverage gaps.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* GRPO no-grad packing: gate verification on real completion rows
Count active rows via create_completion_attention_mask (the same mask the
loss uses) instead of any non-pad token in the packed window. Prompt-only
rows carry prompt-overflow tokens in the window and could otherwise satisfy
the >= 2 verification guard, letting a batch with a single real completion
row cache a trust decision. This matches the gradient path, which already
gates on the completion mask. The same mask is reused for the self-verify
comparison.
* GRPO no-grad packing: gate debug logging on UNSLOTH_ENABLE_LOGGING
Use the shared UNSLOTH_ENABLE_LOGGING global (import_fixes, re-exported by
_utils) instead of a bespoke UNSLOTH_GRPO_SEQ_PACKING_DEBUG env var for the
packing debug prints, matching the rest of the codebase.
* GRPO packing: import UNSLOTH_ENABLE_LOGGING inside the injected logp function
_get_per_token_logps_and_entropies is copied verbatim into the generated GRPO trainer
via inspect.getsource, and that module never imported UNSLOTH_ENABLE_LOGGING, so the
default-on packing verify path raised NameError (and the except handler re-raised it).
Import the flag locally, before the try, so the name is defined in the generated module
too. Drop it from the now-unused module-level import.
* GRPO no-grad packing: harden unsafe-length skip, verify guard, fallback cleanup
Three fixes to the no-grad logp packing path, mirroring the grad path:
- skip the packed forward for known-unsafe lengths by reading unsafe_T and
gating on it before the forward, instead of running the full packed pass and
the result build only to discard them (wastes a pass, can OOM at large T)
- only widen the verified T/seg envelope when >= 2 completion rows actually
exercised cross-sample packing; a < 2 row batch cannot expose leakage, so it
must not extend the trusted shape that later multi-row batches skip verify for
- drop the packed intermediates (hidden/sel/result/ref) before the padded
fallback loop so it does not run with the flattened hidden state still resident
* GRPO no-grad packing: cap the flattened forward at one mini-batch budget
The packed path built a single [1, sum L] forward over every row before any
size check, so a large batch could exceed the memory the padded path bounds
per mini-batch. Gate packing on _pk_T <= _pk_cap (B * seq_len, one padded
mini-batch's token budget); larger batches fall back to the chunked padded
loop.
* GRPO no-grad packing: disable unless unsloth_zoo has the masked-column guard
The packed path leaves masked prompt/pad logprob columns at 0, which only stays
finite if unsloth_zoo grpo_compute_loss zeroes them before exp() (zoo#840). An
older unsloth_zoo without that guard would NaN. Detect the guard once (cached on
the model) via inspect.getsource and gate packing on it, so #6738 is safe with
any unsloth_zoo version and re-enables packing automatically once a guarded zoo
is installed, independent of the pinned lower bound.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* GRPO packing: hoist env gates and zoo-guard detection to one-time module checks
Read UNSLOTH_GRPO_SEQ_PACKING and detect the unsloth_zoo masked-column guard once at
import time (module constants plus RL_PRE_ITEMS for the generated trainer cache) instead
of per call, and drop the in-function UNSLOTH_ENABLE_LOGGING import for a module-top one.
The UNSLOTH_GRPO_SEQ_PACKING_VERIFY force-verify debug knob is commented out, kept in
place for hand re-enable; the first-use and envelope-growth self-verify stays active.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* GRPO packing: cap the flattened forward by the padded chunk rows
B counts chunks at this point, so B * seq_len understated (small runs) or overstated
(large runs) the padded mini-batch token budget; use batch_size * seq_len, the rows the
padded loop actually forwards per chunk.
* Add PrefixGrouper for GRPO: dedup the shared prompt across a group's completions
In GRPO every prompt spawns G=num_generations completions that share the prompt
prefix, so the trunk logprob forward re-encodes that prefix G times. PrefixGrouper
stores the prefix once and concatenates only the G suffixes behind a FlexAttention
shared-prefix mask, cutting the forward from G*(P+R) to P+G*R tokens across both the
no-grad old/ref forwards and the grad logp forward. Default off behind the
UNSLOTH_GRPO_PREFIX_GROUPER env gate, so the gate-unset path is byte-identical to
today. A tok_r auto-gate and a first-use self-verify (fall back and mark the shape
unsafe on mismatch) keep it from ever shipping wrong logprobs silently.
Wired for llama, mistral, qwen3, gemma2, cohere, granite and falcon_h1, plus qwen2
and gemma through the shared LlamaAttention_fast_forward. Stacked on the GRPO
sequence-packing PR (#6738); the grad path lands in a companion unsloth-zoo PR.
Also fixes a latent UNSLOTH_ENABLE_LOGGING NameError in the seq-packing no-grad
verify path by defining the name as a generated-cache pre-item.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* PrefixGrouper: enforce the sliding-window cap, gate softcap models, bound the mask cache
Add a max_segment_cap kwarg to build_group_layout so it falls back when a group's
span (prefix + longest suffix) exceeds the model's local window, and pass the config
sliding_window into the no-grad engage gate the same way the packed _pk guard derives it.
Skip PrefixGrouper entirely for attn_logit_softcapping models, since the FlexAttention
kernel never applies logit softcapping. Bound _BLOCK_MASK_CACHE to a FIFO of 8 so
per-step lengths cannot pin BlockMasks forever, release the PG hidden before the verify
forward, and align the UNSLOTH_ENABLE_LOGGING pre-item truthiness with the canonical form.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* PrefixGrouper: vectorize the real-column scan in build_group_layout
Replace the per-row O(B*L) Python scan of the keep mask with a GPU-derived
contiguous-run fast path (first real column + count per row), keeping the
general scan only as a fallback for non-contiguous rows. Works for both call
sites: the no-grad layout (left-padded prompt + right-padded completion, run
does not start at column 0) and the grad layout (left-packed).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* PrefixGrouper: hoist the gate and kernel imports to one-time module checks, AGPLv3 headers
Read UNSLOTH_GRPO_PREFIX_GROUPER and resolve the prefix_grouper imports once at module
level (source constants plus an RL_PRE_ITEMS entry for the generated trainer cache)
instead of per call, matching the sequence-packing gates. The prefix_grouper env helpers
become one-time module reads with unchanged signatures, and attention_dispatch resolves
the FlexAttention kernel once behind the same gate (lazy fallback kept). The two new
prefix_grouper files move to AGPLv3 headers.
* PrefixGrouper: length-envelope trust and hybrid SSM exclusion
Verified signatures now record (max T, max segment) and re-verify when either grows,
matching the packed path's envelope. Hybrid SSM models (FalconH1 etc.) are excluded at
the gate since only attention gets the shared-prefix isolation, and the FalconH1 wiring
is removed.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* PrefixGrouper: defer the unverified no-grad forward until the packed reference exists
Unverified shapes no longer run the whole-batch shared-prefix forward up front; it now
runs at the verify site, only when the packed path produced a reference. A declined
packed path (budget, window) therefore costs no wasted PG forward per step. Trusted
shapes still run it first to skip the full-row forward, with the same fallback.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* PrefixGrouper: disable under vLLM (fast_inference=True)
With colocated vLLM generation the rollout dominates the GRPO step, so the shared-prefix
training forward saves little end-to-end and its first-use self-verify (which also runs
the full-row path) is net overhead. Gate PG on not use_vllm so it only engages on the raw
transformers path, where the training forward is on the critical path. Packing is unaffected.
* PrefixGrouper: compile the FlexAttention kernel with dynamic shapes
GRPO changes the packed length T almost every batch. With dynamic=False the flex
forward+backward kernel recompiled on every new T (~14s each on a 4B trunk), which
dominated the step and made PG a net loss. dynamic=True compiles once, then reuses the
kernel across all lengths recompile-free (a new shape drops from ~14s to ~1.4ms after a
two-graph warmup). T is still padded to a multiple of 128 for the backward block assertion.
* PrefixGrouper: default on
Enable PrefixGrouper by default (UNSLOTH_GRPO_PREFIX_GROUPER defaults to 1; set 0 to
disable). Still auto-disabled under vLLM (fast_inference=True) and by the arch/softcap/
SSM/tok_r gates, and the first-use self-verify falls back on any mismatch, so this is a
memory-first default on the raw-transformers path with no correctness risk.
* GRPO PrefixGrouper: gate on zoo masked-column guard and exclude MoE
- Require the zoo masked-column guard (zoo#840) before PrefixGrouper can engage.
PG rides the sequence-packing path, so when the first-step self-verify is off the
fast path trusts PG output directly; without the guard those masked columns feed
NaN into the packed loss. Gate PG on the same UNSLOTH_ZOO_HAS_MASKED_COL_GUARD
the packing path already checks.
- Exclude MoE configs (num_experts, num_local_experts, n_routed_experts,
moe_intermediate_size) alongside the hybrid-SSM markers. Only the threaded
attention forwards carry the shared-prefix isolation, so a MoE decoder that does
not forward prefix_seg_info would let suffixes leak across completions.
- Refresh the stale default-off comments now that UNSLOTH_GRPO_PREFIX_GROUPER is
on by default.
* GRPO PrefixGrouper: import chunked_hidden_states_selective_log_softmax
The shared-prefix forward passes chunked_hidden_states_selective_log_softmax
into extract_logps, but the name was only ever provided by the generated
trainer cache (rl.py injects grpo_selective_log_softmax_code), never bound in
this module. Import it from unsloth_zoo.rl_replacements next to its sibling
chunked_selective_log_softmax so the source resolves the name in every scope
(the new _pg_run_forward closure included). No runtime change: the cache still
defines the function via template injection.
* GRPO PrefixGrouper: dropout gate, device-safe layout, Mistral mask skip
Addresses three review findings on the shared-prefix path:
- Skip PrefixGrouper when the model sets a nonzero attention_dropout. The normal
backends apply config.attention_dropout while training (e.g. Granite dense
flash/sdpa/xformers), but the FlexAttention shared-prefix path is deterministic,
so gate PG off for those configs rather than train on mismatched activations.
- Move the shared-prefix mask labels to the consumer (Q) device in get_block_mask
and the target index maps to hidden.device in extract_logps, mirroring the packed
path moving its metadata to the consumer device. Prevents cross-device indexing
when the model is sharded across GPUs.
- Do not synthesize a causal attention_mask in the Mistral forward when
prefix_seg_info is present. On the no-xFormers path that synthetic mask tripped
resolve_prefix_seg_info and forced PG to always fall back to the packed forward.
* GRPO sequence packing: tighten comments
* GRPO PrefixGrouper: tighten comments
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* GRPO PrefixGrouper: persistent disable on runtime failure; build block-mask labels with inference mode disabled
- rl_replacements: on a PG forward exception (FlexAttention/Triton compile failure or OOM), set a model-level _unsloth_prefix_grouper_nograd_disabled flag and consult it in the engage gate, mirroring the seq-packing handler, so a GPU-wide failure is not retried and re-paid every step.
- prefix_grouper_kernel: move the .to(device) label copies inside the inference_mode(False) block so a cross-device (model-parallel shard) first build does not capture inference tensors, which otherwise cannot be saved for backward when the grad training forward reuses the cached BlockMask.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
08e133cd6b