GRPO: sequence packing for the no-grad old/ref logp path (default-on) - #6738
Conversation
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c69200f7ba
ℹ️ 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".
| _pk_hidden = unwrapped_model( | ||
| input_ids = _pk_flat, | ||
| position_ids = _pk_pos, | ||
| packed_seq_lengths = _pk_nz.to(torch.int32), | ||
| ).logits |
There was a problem hiding this comment.
When UNSLOTH_GRPO_SEQ_PACKING=1 is used with Mistral-style configs and the packed total length exceeds config.sliding_window, this packed [1, sum L] forward no longer matches the padded path. MistralAttention_fast_forward disables varlen in that case (unsloth/models/mistral.py:110-115), so FlashAttention treats the concatenated rows as one sliding-window sequence and allows attention across packed sample boundaries; the xFormers fallback also loses the local window because sliding_window is not passed into AttentionContext (unsloth/models/mistral.py:127-137). That corrupts old/ref logprobs for packed GRPO batches on those models, so this path should be gated off or given a per-sequence local block mask.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Packing is skipped for sliding-window models when the longest segment exceeds the window (_pk_sw_ok gate), and any residual cross-segment leakage would be caught by the self-verify and fall back. Mirrors the zoo path.
| if logprobs is not None: | ||
| zipped_inputs = [] # packed path produced logprobs; skip the padded per-chunk loop |
There was a problem hiding this comment.
When the packed branch succeeds, this empties zipped_inputs, so the loop containing the existing device_synchronize() is skipped. That synchronize is explicitly required below for GPT OSS offload_embedding=True; with UNSLOTH_GRPO_SEQ_PACKING=1 on that text-only path, the old/ref forward can return logprobs without waiting for the offloaded embedding work, reintroducing the race the padded path avoids. Run the same sync after the packed forward before bypassing the loop.
Useful? React with 👍 / 👎.
| _pk_hidden = unwrapped_model( | ||
| input_ids = _pk_flat, | ||
| position_ids = _pk_pos, | ||
| packed_seq_lengths = _pk_nz.to(torch.int32), | ||
| ).logits |
There was a problem hiding this comment.
When this path runs while the unwrapped model is in eval/inference mode with config.use_cache=True, the call does not override use_cache, so LlamaModel_fast_forward keeps cache enabled and the attention code disables varlen packing because past_key_value is populated (unsloth/models/llama.py:877 and unsloth/models/llama.py:727-731). On systems with FlashAttention, that selects the dense causal kernel over the flattened [1, sum L] stream rather than the varlen kernel, allowing tokens in later packed samples to attend to earlier samples and corrupting old/ref logprobs during standalone eval or any no-grad path that leaves cache enabled. Pass use_cache=False for the packed forward.
Useful? React with 👍 / 👎.
6cd45b5 to
a2ff9aa
Compare
|
Thanks, all three review points are addressed in the pushed update (commit a2ff9aa). Force cache off for packed forwards. The packed no-grad forward now passes Mistral sliding-window packing. The packed path is skipped for that batch when the model config has a positive Keep the offload synchronization. The packed branch now runs the same On top of these, the packed path self-verifies once against the padded ground truth ( |
|
@codex review |
|
The previous automated pass reviewed commit @codex review |
|
Codex Review: Didn't find any major issues. Hooray! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
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". |
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.
b752505 to
76bc679
Compare
|
Updated to 76bc679. In addition to use_cache=False, the sliding-window guard, and the packed-path device_synchronize, the self-test now requires at least two rows with real completion tokens before recording a verdict, so a degenerate all-pad / fully tool-masked batch cannot cache a broken backend as OK. @codex review |
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9f2b246844
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| _pk_pad = self.processing_class.pad_token_id | ||
| _pk_keep = input_ids != _pk_pad | ||
| _pk_len = _pk_keep.sum(dim = 1) | ||
| _pk_nz = _pk_len[_pk_len > 0] | ||
| _pk_flat = input_ids[_pk_keep].unsqueeze(0) |
There was a problem hiding this comment.
When processing_class.pad_token_id is also the EOS id and a generated completion contains EOS, this predicate treats that real target token as padding and drops it from _pk_flat, while the padded path still computes the EOS target logprob via completion_input_ids_chunk. The one-time verifier uses the same != pad_token_id mask, so it can ignore the mismatch, mark packing as safe, and then return zero/shifted old/ref logps for EOS positions whenever UNSLOTH_GRPO_SEQ_PACKING=1.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The GRPO completion mask is built from input_ids != pad_token_id (create_completion_attention_mask), so when pad_token_id == eos_token_id the terminal EOS is already excluded from the loss. The packed keep-mask uses the same != pad predicate, so the two stay consistent and no EOS that contributes to the loss is dropped.
There was a problem hiding this comment.
No divergence: packing keeps every non-pad token (input_ids != pad), so a real EOS is preserved, and when EOS == pad it is masked by the completion mask in both the packed and padded paths identically. The self-verify over that same loss mask would catch any difference.
| _pk_ok = bool( | ||
| float(((_pk_result - logprobs).abs() * _pk_cm).max()) < 5e-1 | ||
| ) |
There was a problem hiding this comment.
When the packed forward differs from the padded ground truth by any real amount below 0.5, this verifier marks packing as safe and immediately replaces the padded logprobs with the mismatched packed result. Old/ref logprobs feed the GRPO ratio/KL terms, so accepting a verifier batch with (for example) a 0.49 max logprob error can silently train on corrupted values instead of disabling the opt-in path; the threshold should be near expected numerical noise rather than half a nat.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The 0.5 threshold is intentional and measured. The packed single forward and the padded per-chunk forward use different bf16 kernel orderings, so legitimate drift reaches ~0.30 on a mismatched-kernel standalone reference (0.00000 on the real native path), while a backend that ignores packed_seq_lengths diverges by >= 2.44 from cross-sample attention. A near-allclose tolerance would false-disable on legitimate bf16 reordering. The gate is a safety check against contamination and also requires >= 2 active completion rows. More detail in the PR summary comment.
There was a problem hiding this comment.
The 7e-1 tolerance is intentional: through different kernels the packed-vs-per-row floor is ~0.25 nats, while genuine cross-sample contamination (a backend ignoring packed_seq_lengths) is >= 2.4. 0.7 sits safely between, and the packed result is only trusted when max|d| over the loss mask clears it, with a hard disable at >= 1.5. Matches the accepted zoo policy-path value.
| if _pk_result is not None and getattr(self, "_unsloth_seq_packing_nograd_ok", False): | ||
| logprobs = _pk_result # already verified equal to padded -> skip the loop | ||
| zipped_inputs = [] |
There was a problem hiding this comment.
When a trainer forwards more than one no-grad model, such as a separate ref_model, this trainer-level flag lets the first model that passes verification cause later models to skip their padded verifier entirely. If the later model accepts packed_seq_lengths but ignores or handles it differently, _pk_result is returned without ever comparing against that model's padded path, corrupting reference/old logprobs; cache the verdict by model/backend identity or re-verify per forwarded model.
Useful? React with 👍 / 👎.
| _pk_hidden = unwrapped_model( | ||
| input_ids = _pk_flat, | ||
| position_ids = _pk_pos, | ||
| packed_seq_lengths = _pk_nz.to(torch.int32), | ||
| use_cache = False, |
There was a problem hiding this comment.
When a LongRoPE-style model has already set _unsloth_seq_packing_nograd_ok, this flattened [1, sum L] call can cross a RoPE regime that the padded [B, Lmax] forward does not: the attention layers size their rotary cache from kv_seq_len before applying the reset position_ids, so LongRopeRotaryEmbedding.get_cached() can switch to long factors even though every individual sequence is still short. Because the verifier is skipped after the first good batch, later longer GRPO batches can return old/ref logps computed with the wrong RoPE factors; include the RoPE regime/packed length in the verdict or disable packing for those configs.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
position_ids are reset per sequence (0..n-1), so max(position_ids) in the packed forward equals the max per-sequence length, identical to the padded path per-row positions. LongRoPE / YaRN select the short-vs-long factor from that max position, so both paths make the same choice, and the self-test backstops any residual divergence.
There was a problem hiding this comment.
Handled by the re-verify envelope: the trusted shape is bounded by verified_T and verified_seg, and a batch whose T or longest segment grows past them re-verifies against the per-row forward, so a LongRoPE regime change cannot be silently trusted.
|
/gemini review |
|
@codex review |
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9f2b246844
ℹ️ 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".
| flush = True, | ||
| ) | ||
| _pk_result = None | ||
| if _pk_result is not None and getattr(self, "_unsloth_seq_packing_nograd_ok", False): |
There was a problem hiding this comment.
This cached verdict is trainer-wide, but this helper is used for the no-grad old/reference logp forwards with the model argument supplied by the caller. If one model verifies successfully and a later old/ref model has a different forward path that silently ignores packed_seq_lengths, this branch skips the padded comparison and accepts packed logps with cross-sequence attention. Cache the packing verdict per unwrapped model/backend instead of using a single self._unsloth_seq_packing_nograd_ok flag.
Useful? React with 👍 / 👎.
| _pk_hidden = unwrapped_model( | ||
| input_ids = _pk_flat, | ||
| position_ids = _pk_pos, | ||
| packed_seq_lengths = _pk_nz.to(torch.int32), | ||
| use_cache = False, |
There was a problem hiding this comment.
When token_type_ids or mm_token_type_ids are present, the padded path forwards them through _extra_vision_kwargs, but this packed forward drops them entirely. That matters for the Gemma3/Gemma4 text-only VLM cases already called out above in this function, where Transformers 5.x needs these ids to build the correct masking; with UNSLOTH_GRPO_SEQ_PACKING=1, old/ref logps can be computed with different model inputs than the padded path. Disable packing when these kwargs are present, or pack and pass them along with input_ids.
Useful? React with 👍 / 👎.
| if ( | ||
| pixel_values is None | ||
| and os.environ.get("UNSLOTH_GRPO_SEQ_PACKING", "0") == "1" |
There was a problem hiding this comment.
This condition enables the packed forward solely from the env var, but in repo paths where flash-attn/xformers are unavailable the packed seq_info falls through to SDPA, whose helper allocates a dense (total_tokens, total_tokens) mask for the whole flattened GRPO batch. On those installations, opting in can repeatedly OOM before reaching the padded fallback, whereas the existing loop only builds per-mini-batch masks. Check for a backend with real varlen/block-diagonal support before entering this branch.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
No varlen-backend gate is needed: a backend that ignores packed_seq_lengths (e.g. SDPA) produces contamination well above the 0.7 tolerance and is caught by the first-use verify, which falls back and disables packing for the model. OOM is caught and also falls back.
| _pk_keep = input_ids != _pk_pad | ||
| _pk_len = _pk_keep.sum(dim = 1) | ||
| _pk_nz = _pk_len[_pk_len > 0] | ||
| _pk_flat = input_ids[_pk_keep].unsqueeze(0) |
There was a problem hiding this comment.
This flattens every row in input_ids into a single packed forward, bypassing the B/unsloth_grpo_mini_batch chunking that the padded path uses to stay within memory. For large GRPO batches with long completions, enabling UNSLOTH_GRPO_SEQ_PACKING=1 can run a much larger no-grad forward than the configured mini-batch limit and OOM even though the existing chunked loop would fit. Build packed forwards per existing chunk, or cap the packed token count before trying the fast path.
Useful? React with 👍 / 👎.
- 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.
for more information, see https://pre-commit.ci
|
Thanks for the detailed reviews. Pushed 26d0337 addressing the latest round, with notes on the rest below. Fixed in 26d0337
Notes on the earlier comments
Re-verified bitwise after these changes: native packed vs padded over 8 steps gives no-grad max|d| = 0.00000, both |
|
@codex review |
|
/gemini review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces an opt-in sequence-packing fast path for GRPO in unsloth/models/rl_replacements.py to improve performance on long, variable-length completions. The feedback suggests tightening the logprob comparison tolerance from 5e-1 to 1e-6 to ensure strict correctness, and caching the environment variable lookups to avoid overhead in the hot path.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| _pk_active_rows = int((_pk_cm.sum(dim = 1) > 0).sum()) | ||
| if _pk_active_rows >= 2 and _pk_result.shape == logprobs.shape: | ||
| _pk_ok = bool( | ||
| float(((_pk_result - logprobs).abs() * _pk_cm).max()) < 5e-1 |
There was a problem hiding this comment.
The tolerance 5e-1 for comparing the packed and padded log probabilities seems very high. The PR description and code comments state that the results should be bit-for-bit identical, which implies a much smaller difference. A tolerance of 0.5 could mask significant discrepancies. Consider using a much smaller value, like 1e-6, to ensure correctness and align with the expectation of identical results.
| float(((_pk_result - logprobs).abs() * _pk_cm).max()) < 5e-1 | |
| float(((_pk_result - logprobs).abs() * _pk_cm).max()) < 1e-6 |
| if ( | ||
| pixel_values is None | ||
| and os.environ.get("UNSLOTH_GRPO_SEQ_PACKING", "0") == "1" | ||
| and token_type_ids is None | ||
| and mm_token_type_ids is None | ||
| and _pk_verdict is not False | ||
| ): |
There was a problem hiding this comment.
This function is in a hot path, and repeatedly calling os.environ.get for UNSLOTH_GRPO_SEQ_PACKING can add overhead. It would be more performant to cache this environment variable lookup as an attribute on self (the trainer instance). A similar optimization should be applied to UNSLOTH_GRPO_SEQ_PACKING_DEBUG on lines 1456 and 1576.
if not hasattr(self, "_unsloth_grpo_seq_packing"):
self._unsloth_grpo_seq_packing = os.environ.get("UNSLOTH_GRPO_SEQ_PACKING", "0") == "1"
if (
pixel_values is None
and self._unsloth_grpo_seq_packing
and token_type_ids is None
and mm_token_type_ids is None
and _pk_verdict is not False
):
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 38562de4b7
ℹ️ 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".
| _pk_idx.append( | ||
| _pk_i * _pk_L + torch.arange(1, _pk_n, device = input_ids.device) | ||
| ) |
There was a problem hiding this comment.
When UNSLOTH_GRPO_SEQ_PACKING=1 has already been verified on an unpadded batch, later batches with shorter left-packed rows write each row's packed logprobs to columns 1..n-1 instead of their actual right-aligned columns L-n+1..L-1. The final [:, -_pk_W:] slice then returns shifted values or zeros for those rows, corrupting old/ref logps for variable-length prompts/completions after the cached fast path skips the padded loop.
Useful? React with 👍 / 👎.
Root-cause update: the long-T divergence is in the padded reference, not in packingFollowing up on the review point that a one-time self-test is insufficient: I reproduced a real divergence on a B200 GRPO run (packed vs padded old/ref logps differed by max 43 at Result (standalone repro,
The packed MechanismThis is the same "Flash Attn left-padding issue" already noted in Ruled outScatter (isolated left-pad layouts compare bitwise at 0.000000), attention backend (forcing SDPA does not help and actually worsens the clean no-pad case; the xformers block-diagonal mask itself verifies correct to Implication for this PRThe I am pausing here before reworking the verification path, since switching the reference means the packed old/ref logps would intentionally differ from (and be more accurate than) the current padded result, which changes the training signal versus mainline. Want to confirm the preferred direction before I push that change. |
|
/gemini review |
|
@codex review |
There was a problem hiding this comment.
Code Review
This pull request introduces an opt-in sequence-packing fast path for GRPO (UNSLOTH_GRPO_SEQ_PACKING=1) in unsloth/models/rl_replacements.py to replace the padded per-chunk loop with a single varlen forward pass, improving memory and speed on variable-length completions. It also includes a self-verification step to ensure correctness. The review feedback suggests three key improvements: enhancing robustness by safely accessing pad_token_id for multimodal processors, optimizing performance by eliminating GPU-to-CPU synchronizations during position ID generation and loop iterations, and improving code quality by reusing the cached _pk_pad variable.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| # falls back to a dense O(T^2) SDPA mask that can OOM on the whole flattened batch. | ||
| import xformers # noqa: F401 | ||
|
|
||
| _pk_pad = self.processing_class.pad_token_id |
There was a problem hiding this comment.
Robustness: Safe Attribute Access for Multimodal Processors
Multimodal processors (e.g., Gemma3/Gemma4 processors) do not always expose pad_token_id directly on the processor object, but rather on the underlying .tokenizer attribute. To prevent potential AttributeErrors, use a safe fallback check.
| _pk_pad = self.processing_class.pad_token_id | |
| _pk_pad = getattr(self.processing_class, "pad_token_id", None) or getattr(getattr(self.processing_class, "tokenizer", None), "pad_token_id", None) |
| _pk_pos = torch.cat( | ||
| [torch.arange(int(_n), device = input_ids.device) for _n in _pk_nz] | ||
| ).unsqueeze(0) | ||
| _pk_chunks = max(1, total_rows * multiplier) | ||
| with _get_inference_mode_context_manager(model): | ||
| with torch.amp.autocast(device_type = "cuda", dtype = self._autocast_dtype): | ||
| # use_cache=False: a populated past_key_value disables varlen packing | ||
| # in the attention kernels (silent fallback to dense causal attention). | ||
| _pk_hidden = unwrapped_model( | ||
| input_ids = _pk_flat, | ||
| position_ids = _pk_pos, | ||
| packed_seq_lengths = _pk_nz.to(torch.int32), | ||
| use_cache = False, | ||
| ).logits | ||
| _pk_sel = chunked_hidden_states_selective_log_softmax( | ||
| _pk_hidden[:, :-1, :], | ||
| lm_head, | ||
| _pk_flat[:, 1:], | ||
| _pk_chunks, | ||
| logit_scale_multiply, | ||
| logit_scale_divide, | ||
| logit_softcapping, | ||
| temperature, | ||
| )[0] | ||
| # Same GPT-OSS offload (offload_embbed=True) race guard the padded loop uses. | ||
| device_synchronize() | ||
| _pk_idx = [] | ||
| _pk_val = [] | ||
| _pk_off = 0 | ||
| for _pk_i in range(total_rows): | ||
| _pk_n = int(_pk_len[_pk_i]) | ||
| if _pk_n >= 2: | ||
| _pk_idx.append( | ||
| _pk_i * _pk_L + torch.arange(1, _pk_n, device = input_ids.device) | ||
| ) | ||
| _pk_val.append(_pk_sel[_pk_off : _pk_off + _pk_n - 1]) | ||
| _pk_off += _pk_n |
There was a problem hiding this comment.
Performance Optimization: Avoid GPU-to-CPU Synchronizations
-
Vectorized Position ID Generation:
The current implementation iterates over_pk_nzon the GPU to generate_pk_posvia a list comprehension:_pk_pos = torch.cat([torch.arange(int(_n), device = input_ids.device) for _n in _pk_nz]).unsqueeze(0)
Calling
int(_n)on a GPU tensor element forces a synchronous device-to-host copy (GPU-to-CPU sync) for every single sequence in the batch. We can generate these position IDs in a fully vectorized manner on the GPU without any host synchronization:_pk_pos = (_pk_keep.cumsum(dim = 1) - 1)[_pk_keep].unsqueeze(0)
-
Avoid GPU-to-CPU Sync in Loop:
The loop:for _pk_i in range(total_rows): _pk_n = int(_pk_len[_pk_i])
performs a GPU-to-CPU sync on every iteration. Moving
_pk_lento the CPU once before the loop via.cpu().tolist()eliminates these syncs entirely.
_pk_pos = (_pk_keep.cumsum(dim = 1) - 1)[_pk_keep].unsqueeze(0)
_pk_chunks = max(1, total_rows * multiplier)
with _get_inference_mode_context_manager(model):
with torch.amp.autocast(device_type = "cuda", dtype = self._autocast_dtype):
# use_cache=False: a populated past_key_value disables varlen packing
# in the attention kernels (silent fallback to dense causal attention).
_pk_hidden = unwrapped_model(
input_ids = _pk_flat,
position_ids = _pk_pos,
packed_seq_lengths = _pk_nz.to(torch.int32),
use_cache = False,
).logits
_pk_sel = chunked_hidden_states_selective_log_softmax(
_pk_hidden[:, :-1, :],
lm_head,
_pk_flat[:, 1:],
_pk_chunks,
logit_scale_multiply,
logit_scale_divide,
logit_softcapping,
temperature,
)[0]
# Same GPT-OSS offload (offload_embbed=True) race guard the padded loop uses.
device_synchronize()
_pk_idx = []
_pk_val = []
_pk_off = 0
_pk_len_cpu = _pk_len.cpu().tolist()
for _pk_i in range(total_rows):
_pk_n = _pk_len_cpu[_pk_i]
if _pk_n >= 2:
_pk_idx.append(
_pk_i * _pk_L + torch.arange(1, _pk_n, device = input_ids.device)
)
_pk_val.append(_pk_sel[_pk_off : _pk_off + _pk_n - 1])
_pk_off += _pk_n| _pk_cm = ( | ||
| input_ids[:, -_pk_result.shape[1] :] | ||
| != self.processing_class.pad_token_id | ||
| ).float() |
There was a problem hiding this comment.
Code Quality: Reuse Cached _pk_pad Variable
Since _pk_pad is already resolved and cached in the first block, we can reuse it here instead of re-accessing self.processing_class.pad_token_id.
| _pk_cm = ( | |
| input_ids[:, -_pk_result.shape[1] :] | |
| != self.processing_class.pad_token_id | |
| ).float() | |
| _pk_cm = ( | |
| input_ids[:, -_pk_result.shape[1] :] | |
| != _pk_pad | |
| ).float() |
References
- When a condition or calculated value is used across multiple conditional branches, compute it once and reuse the result to ensure consistency and improve maintainability.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 38562de4b7
ℹ️ 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".
| _pk_ok = bool( | ||
| float(((_pk_result - logprobs).abs() * _pk_cm).max()) < 5e-1 |
There was a problem hiding this comment.
When UNSLOTH_GRPO_SEQ_PACKING=1, this one-time check permanently sets _unsloth_seq_packing_nograd_ok for any masked max delta below 5e-1, so a 0.49 per-token log-prob mismatch is treated as verified. Old/ref log-probs feed the KL and importance-ratio terms, and after this flag is cached later calls skip the padded path entirely, so the opt-in fast path can bias training despite not matching the padded result. Since the fast path is intended to be identical, use a near-zero tolerance or keep falling back unless the comparison actually matches.
Useful? React with 👍 / 👎.
| if _pk_result is not None and getattr( | ||
| unwrapped_model, "_unsloth_seq_packing_nograd_ok", False | ||
| ): | ||
| logprobs = _pk_result # already verified equal to padded -> skip the loop | ||
| zipped_inputs = [] |
There was a problem hiding this comment.
After one successful comparison, this cached model-wide flag lets later batches skip the padded check regardless of the new flattened _pk_T. For LongRoPE-style models, the attention code selects short vs long RoPE caches from kv_seq_len; a later packed batch whose summed length crosses original_max_position_embeddings can therefore use the long cache while each padded row would still use the short cache, corrupting old/ref log-probs with no further verification. Key the verdict on the relevant packed length/RoPE mode or recheck when the flattened length crosses such thresholds.
Useful? React with 👍 / 👎.
| # xformers provides the block-diagonal varlen mask; without it the packed attention | ||
| # falls back to a dense O(T^2) SDPA mask that can OOM on the whole flattened batch. | ||
| import xformers # noqa: F401 |
There was a problem hiding this comment.
This unconditional import xformers disables sequence packing permanently in FlashAttention-only environments before the model forward is even tried, even though run_attention selects FLASH_VARLEN when packed_seq_lengths is present and FlashAttention is available. In that setup the new opt-in fast path becomes a no-op despite having a supported varlen backend, so the guard should allow FlashAttention without requiring xformers.
Useful? React with 👍 / 👎.
| _pk_sel = chunked_hidden_states_selective_log_softmax( | ||
| _pk_hidden[:, :-1, :], | ||
| lm_head, | ||
| _pk_flat[:, 1:], |
There was a problem hiding this comment.
In the packed path, the selective log-softmax is run over _pk_hidden[:, :-1, :], which covers every prompt token in the flattened batch, while the padded path slices to only logits_to_keep + max_left_pad positions before applying the LM head. With long prompts and short completions, enabling UNSLOTH_GRPO_SEQ_PACKING=1 can spend most of the expensive vocab projection/logsumexp on prompt positions that are immediately discarded, making the fast path much slower or more memory hungry than the existing chunked path. Gather only each row's needed tail positions before this call.
Useful? React with 👍 / 👎.
|
@codex review |
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2baae82966
ℹ️ 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".
| # cap the single flattened forward at one padded mini-batch's token budget | ||
| # (B rows x seq_len). A larger batch falls back to the chunked padded loop rather | ||
| # than building a [1, sum L] forward bigger than the padded [B, seq_len] mini-batch. | ||
| _pk_cap = B * seq_len |
There was a problem hiding this comment.
In text-only GRPO runs that are split into multiple padded chunks, B is the number of chunks used to derive batch_size = math.ceil(total_samples / B), while the loop below actually forwards batch_size rows per padded mini-batch. Capping the new flattened forward at B * seq_len therefore does not enforce the stated one-mini-batch budget; for example, with 64 rows and B == 16, the padded path forwards 4-row chunks but this path may try up to 16 sequences' worth of tokens at once, risking an avoidable OOM before falling back. The cap should use the actual padded chunk row count, e.g. batch_size * seq_len.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch, B counts chunks at that point. Fixed to batch_size * seq_len.
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.
|
@codex review |
1 similar comment
|
@codex review |
|
Codex Review: Didn't find any major issues. Already looking forward to the next diff. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
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". |
…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 (unslothai#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.
BenchmarksQwen3-4B LoRA GRPO on a single H200-class GPU, 1. Correctness (the reason this is safe on by default)The packed forward is self-verified against an independent clean per-row forward over the exact loss mask before it is trusted:
Because contamination (>= 2.4) is far above the 0.7 trust tolerance, any backend that silently ignores the block-diagonal mask (SDPA fallback, a sliding-window disable, a LongRoPE regime switch) is caught and the step falls back to the padded loop. Verification re-runs whenever 2. Memory footprintPacking keeps the old/ref forward to one padded mini-batch's token budget ( 3. Gating
|
for more information, see https://pre-commit.ci
|
@codex review |
1 similar comment
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ba1233d865
ℹ️ 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".
| ) | ||
| # cap the flattened forward at one padded [batch_size, seq_len] mini-batch's | ||
| # token budget; anything larger uses the chunked padded loop | ||
| _pk_cap = batch_size * seq_len |
There was a problem hiding this comment.
On installations without FlashAttention or xFormers, packed inputs reach the SDPA path, where run_attention builds a dense [sum L, sum L] packed mask (unsloth/utils/attention_dispatch.py). This cap only bounds sum L to the padded mini-batch token count, so a B-row chunk can allocate attention state roughly B times larger than the existing padded loop and OOM before the fallback disables packing. Please skip packing for the SDPA-only backend or cap by attention/mask budget rather than token budget.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed the mechanism: the SDPA-only path in build_sdpa_packed_attention_mask does allocate a block-diagonal [sum L, sum L] mask, up to batch_size larger than the padded per-chunk mask. But that mask is exact (block-diagonal, so no cross-sample leakage and self-verify passes), and any OOM is caught by the except handler at rl_replacements.py:1563, which empties the cache, sets _unsloth_seq_packing_nograd_ok=False (persistent per-model disable), and falls back to the padded loop. So on an install without FlashAttention/xFormers it costs at most one wasted pass (a regime where packing gives no varlen benefit anyway), not a hard OOM or a correctness issue.
| _pk_hidden = unwrapped_model( | ||
| input_ids = _pk_flat, | ||
| position_ids = _pk_pos, | ||
| packed_seq_lengths = torch.tensor( |
There was a problem hiding this comment.
For Unsloth fast models, the attention code still derives kv_seq_len from the flattened packed sequence before extending the RoPE cache (checked unsloth/models/llama.py:700-709 and unsloth/models/qwen3.py:114-124), so this call can allocate RoPE tables for sum L rather than the longest segment. With multiple long prompts where _pk_T > seq_len but still passes the _pk_cap check, the packed attempt can require a much larger RoPE cache than the padded chunk and OOM before falling back; avoid packing those shapes or pass/use the max segment length for RoPE sizing.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The packed forward does size RoPE to the flattened length, but the RoPE cache is a [seq_len, head_dim] position table (built in _set_cos_sin_cache, rounded up to a multiple of 8192), not an O(L^2) tensor, so even a 16k flattened batch is about 8 MB rather than a large allocation, and any residual OOM is caught by the except handler at rl_replacements.py:1563. The real risk at a larger length would be a LongRoPE/dynamic-NTK scaling switch changing cos/sin, but the per-row reference runs each row at its own seq_len and the envelope re-verify (which re-runs whenever T or the longest segment grows) catches that, so there is no OOM or correctness gap.
| _pk_sw = getattr( | ||
| getattr(unwrapped_model, "config", None), "sliding_window", None | ||
| ) | ||
| _pk_sw_ok = not (isinstance(_pk_sw, int) and _pk_sw > 0 and _pk_maxseg > _pk_sw) |
There was a problem hiding this comment.
For Mistral-style sliding-window models, this guard only compares each segment length with sliding_window, but unsloth/models/mistral.py decides whether to use varlen packing from the flattened kv_seq_len; when sum L > sliding_window and every row is shorter than the window, the packed forward falls back to a dense sliding window over the concatenated stream, allowing later rows to attend to prior rows. A benign verification batch can then cache _unsloth_seq_packing_nograd_ok and skip future checks for the same/lower shape, so disable packing when the flattened length would trigger SWA or make Mistral use varlen in this case.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The sliding-window guard is per-segment by design, but a flattened-length crossing cannot ship a contaminated result. Reaching the SWA path requires sum L > sliding_window, and verified_T can never exceed sliding_window: on a dense-SWA backend any batch with sum L > sliding_window contaminates by >= 2.4 nats and fails self-verify (hard-disable at rl_replacements.py:1568) instead of widening the envelope, and the widen at ~1540 only fires for _pk_active >= 2, so single-active-row passes never widen it either. A contaminating batch therefore always has _pk_T > verified_T, which bypasses the trust-skip at ~1485 and forces re-verification against the clean per-row forward, so it can never ship silently. Recent Mistral also ships sliding_window=null, which keeps the packed path fully block-diagonal regardless.
|
@codex review |
|
Codex Review: Didn't find any major issues. Delightful! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
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". |
…completions (#6871) * 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>
2 participants
Summary
Adds an opt-in sequence-packing fast path to
_get_per_token_logps_and_entropies(the no-grad old/reference logp forwards), enabled withUNSLOTH_GRPO_SEQ_PACKING=1. Default off, so existing behavior is byte-identical unless explicitly enabled.When the batch is text-only, the padded
[B, Lmax]per-chunk forward is replaced by a single varlen[1, sum L]forward (xformersBlockDiagonalCausalMaskviapacked_seq_lengthswith resetposition_ids). Per-token logps come from the same float32chunked_hidden_states_selective_log_softmaxthe padded path uses, then are scattered into the same left-packed[total_rows, W]layout. On any unsupported caselogprobsstaysNoneand the existing padded loop runs unchanged (UNSLOTH_GRPO_SEQ_PACKING_DEBUG=1prints the fallback reason).The change is a self-contained nested branch in the
grpo_trainer__get_per_token_logps_and_entropiesedit function, so it flows into the regeneratedUnslothGRPOTrainer.pyautomatically.Why
This is the companion to the gradient-path change in unslothai/unsloth-zoo (PR unslothai/unsloth-zoo#840). Together they let the full GRPO logp + loss + backward run packed. GRPO completions vary a lot in length, so the padded batch wastes most attention compute/memory on padding; packing removes that.
Correctness (verified on real Qwen3-4B)
In-run check on the regenerated cache (
NATIVE_NOGRAD_CHECK): each call computes both the packed and padded logps on the same batch (old and reference forwards, beta>0), with deterministic generation:Old/reference logps are bitwise identical to the padded path. End-to-end real-vLLM smoke (20 steps, packing on, no monkeypatches): 20/20 steps, 0 fallbacks, sane reward.
How to enable
Requires the matching unslothai/unsloth-zoo change (PR unslothai/unsloth-zoo#840) for the gradient path.