2 commits

Author	SHA1	Message	Date
Daniel Han	baa96a7160	Stop the GRPO hidden-states wrapper paying for logits it discards (#8576) ... * Stop the GRPO hidden-states wrapper paying for logits it discards `Kaggle-Muse_Glimmer_(30B)-GRPO` dies on a 2 x T4 kernel at the first training step: accelerate/hooks.py:429 AlignDevicesHook.post_forward -> send_to_device(output, self.input_device) OutOfMemoryError: Tried to allocate 1002.00 MiB. GPU 0 has 14.56 GiB capacity, of which 768.81 MiB is free. It was reported before that as `CUDA error: an illegal memory access was encountered` inside `transformers.masking_utils.fast_all`, which is a `tensor.sum()` and therefore the first synchronising op after the real fault. Re-running with CUDA_LAUNCH_BLOCKING=1 resolved it to the OOM above. Two costs meet in that copy, and this wrapper creates both. UnslothEfficientGRPO never sees a logits tensor: it takes per-token logps plus `lm_head` and chunks the projection itself, and this wrapper exists to hand it hidden states in place of logits. But it forwarded the caller's `logits_to_keep` unchanged, and the GRPO trainer does not pass one: outputs = unwrapped_model( input_ids = input_ids_chunk, attention_mask = attention_mask_chunk, ... ) logits_chunk = outputs.logits del outputs transformers reads a missing or zero value as `slice(-0, None)`, which is `slice(0, None)`: every position. So the model projected the whole prompt and completion over a 202048-wide vocabulary, softcapped the result twice, and the wrapper then overwrote it with hidden states. `output_hidden_states = True` also returns every layer. Only `[-1]` is read and the rest stayed attached to the output. On one card neither cost is visible; `del outputs` frees both a line later. Under an accelerate layer-split dispatch, `io_same_device` walks the whole returned object and copies every tensor in it to the input device first, so both cross the bus and the first card runs out. ## The repair `_minimise_logits_kwarg` pins the limit to 1 when we are going to replace the logits anyway. 1 rather than 0 because 0 means "all of them". It prefers `logits_to_keep`, falls back to the legacy `num_logits_to_keep`, and does nothing when the forward takes neither, when the caller passed the value positionally (passing it again by keyword is a TypeError), or when `labels` is present, since a model computing its own loss needs real logits. If a forward advertises the parameter and then rejects the value, the call is retried without it rather than losing hidden states over an optimisation. `_drop_spare_hidden_states` clears the layers nobody reads. Note that `outputs.hidden_states = None` does NOT do this. ModelOutput.__setattr__ is if name in field_names and value is not None: super().__setitem__(name, value) super().__setattr__(name, value) so assigning None updates the attribute and leaves the mapping entry holding the full tuple, and `__delitem__`, `pop`, `update` and `setdefault` all raise. Anything that walks the object as a mapping -- `send_to_device`, which is the one that matters -- still sees and copies all of it. Writing through `OrderedDict.__setitem__` is what actually clears it. ## Measured Real forward, `unsloth/Qwen2.5-0.5B-Instruct`, 401 positions, 151936 vocab, transformers 4.57.6: lm_head rows peak MiB hidden_states entry returned shape before 401 1094.1 tuple (1, 401, 896) after 1 989.6 None (1, 401, 896) The returned tensor is identical, so nothing downstream sees a difference. ## Scope The wrapper is only installed for models that do not natively honour UNSLOTH_RETURN_HIDDEN_STATES; `_model_supports_unsloth_return_hidden_states` short-circuits the ones unsloth patches itself, so llama, mistral and the rest of the fast paths are untouched. The `outputs.hidden_states` reads in `llama.py` and `mistral.py` are inside their own forwards, on the inner model's output, and run before this wrapper's return. ## Tests `tests/test_grpo_hidden_states_logits_cost.py`, 23 tests. The fake output object reproduces ModelOutput's assignment semantics exactly, and one test asserts the trap still exists so the rest cannot pass vacuously. Both spellings of the kwarg, the positional case, the `labels` case, var-keyword forwards, the reject-and-retry path, and an unrelated TypeError still propagating. Mutation-checked: dropping the minimisation turns 3 red, and replacing the mapping write with the naive `= None` turns the layer-drop test red. 1019 GRPO and hidden-state tests pass together. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep the raw-logits fallback whole when the logits limiter is in play The wrapper now asks for one logit position because it overwrites them with hidden states. Three paths still hand those logits back: - a forward that refuses the limiter and then also refuses output_hidden_states raised its second TypeError outside the handler, where the first would have fallen back to raw logits. Route the retry through the same fallback. - a forward that accepts the limiter but returns no hidden states returned the one-position tensor as the result. GRPO drops the last position and slices the completion window out of the rest, so it had nothing left. Re-run on the caller's own arguments instead. - labels supplied positionally never reached the keyword lookup that keeps the logits intact, so the model computed its own loss over one position. Read them off the bound arguments. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Re-run the absent-hidden-states fallback on the shape that just worked The caller's own kwargs could carry a value already bound positionally; the minimised kwargs are the ones the forward has just accepted, so put the caller's limit back into those and reuse them. * Do not reach for the legacy kwarg when a width is already positional Two things, both from measuring the transformers range this has to work across rather than assuming it. The comment was wrong. It said 4.57.6 spells the limit num_logits_to_keep. inspect.signature on LlamaForCausalLM.forward in real 4.57.6, 5.0.0 and 5.15.0 venvs says all three declare logits_to_keep and none declares the other name. So the legacy branch is not there for transformers; it is there for us. unsloth/models/llama.py and mistral.py patch in forwards declaring both, and unsloth/models/vision.py probes the old name first because some VLM stacks still carry only it. The branch stays; the reason is now stated correctly. The behaviour change: when the modern name is already bound positionally, give up instead of falling through to the legacy one. On a forward that declares only the modern name and sinks the rest into **kwargs, the legacy name is accepted and ignored -- no logits saved -- and returning a non-None name arms the absent-hidden-states re-run, so we would buy a second full forward for nothing. A caller that bound a width positionally has said what it wants; leave it alone. New test covers exactly that shape. 27 tests pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: danielhanchen <unslothshared@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>	2026-08-13 05:15:07 -07:00
Daniel Han	42ed15f4cc	GRPO: dispatch on width at the remaining lm_head matmul call sites (#8204) ... * GRPO: dispatch on width at the remaining lm_head matmul call sites * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO: read an explicit hidden-states signal when the widths cannot decide The dispatch guard compared the forward's last dim against lm_head.shape[1], which is the hidden size. On a model whose vocab_size equals its hidden size, real logits satisfy that test and were sent through chunked_hidden_states_selective_log_softmax, applying the lm_head a second time and silently corrupting the log probabilities. The packed path's per-row verifier made the same misclassification, so it compared corrupted against corrupted and accepted the result. Route all four call sites through _unsloth_grpo_returns_hidden_states, which reads an explicit signal that the forward honoured UNSLOTH_RETURN_HIDDEN_STATES: __UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__ written by the zoo compiler, or the _unsloth_grpo_hidden_states_forward_wrapped pair set by the rl.py fallback wrapper. The width comparison stays: it is decisive whenever vocab_size differs from hidden_size, and the signal is only consulted for the square case the shape cannot answer, so an unsloth_zoo old enough to write no marker keeps today's behaviour. * GRPO: propagate hidden-state signal to gradient dispatches * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Track GRPO hidden-state degradation per call, and reject a half-patched zoo The GRPO width dispatch reads two things the code assumed but did not hold. _warn_grpo_hidden_states_fallback_once only ever set _unsloth_grpo_hidden_states_warning_issued, so the flag the dispatch reads after a forward meant "ever degraded", not "degraded on this call". Degradation is per call: a forward that splats **kwargs into a sub-module only some inputs reach rejects the hidden-state request on those batches and honours it on the rest. With vocab_size == hidden_size the width test cannot correct that, so one degraded batch sent every later hidden-state tensor to the raw-logits helper, skipping the lm_head matmul. Record the outcome of each call in _unsloth_grpo_hidden_states_degraded and keep the warning flag for warn-once logging only; the signal reader falls back to the old flag when the attribute is absent, so a stale generated trainer keeps working. The fallback's TypeError retry also could not work: _drop_forward_kwargs_ consumed_positionally hands the caller's dict straight back when there is nothing to drop, which every GRPO call site hits since they pass everything by keyword, so adding output_hidden_states/return_dict poisoned the caller's kwargs and the retry re-sent exactly what the model had just rejected. Copy before mutating. Finally, the source patch over zoo's gradient dispatches only failed when nothing matched. A zoo that respells some of its dispatch sites still leaves one the pattern recognises, which was enough to suppress the compatibility error while the respelled sites kept deciding on width alone. Count the branch headers that decide off an lm_head dimension and require that none survive the substitution. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Daniel Han <unslothshared@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>	2026-08-11 10:44:06 -07:00