mirror of
https://github.com/unslothai/unsloth.git
synced 2026-08-16 12:24:06 +00:00
baa96a7160
* 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>
77 lines
2.6 KiB
Python
77 lines
2.6 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
|
|
|
"""Load the GRPO hidden-states forward wrapper out of ``unsloth/models/rl.py``.
|
|
|
|
Same trick as ``_grpo_dispatch_source``: lift the module-level defs with ``ast``
|
|
rather than importing ``unsloth``, so the tests stay CPU-only and import-free
|
|
while still tracking the shipped code.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ast
|
|
import collections
|
|
import inspect
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
|
|
|
|
SOURCE_PATH = Path(__file__).resolve().parents[1] / "unsloth" / "models" / "rl.py"
|
|
|
|
WRAPPER_NAMES = (
|
|
"_module_returns_logits",
|
|
"_grpo_hidden_states_wrap_target",
|
|
"_model_supports_unsloth_return_hidden_states",
|
|
"_drop_forward_kwargs_consumed_positionally",
|
|
"_get_num_logits_to_keep",
|
|
"_warn_grpo_hidden_states_fallback_once",
|
|
"_note_grpo_hidden_states_success",
|
|
"_replace_outputs_logits",
|
|
"_minimise_logits_kwarg",
|
|
"_drop_spare_hidden_states",
|
|
"_install_grpo_hidden_states_forward_wrapper",
|
|
)
|
|
|
|
# present only once the per-call degradation fix has landed
|
|
OPTIONAL_NAMES = ("_note_grpo_hidden_states_success",)
|
|
|
|
CONSTANT_NAMES = (
|
|
"_UNSLOTH_RETURN_HIDDEN_STATES_SUPPORT_MARKER",
|
|
"_UNSLOTH_GRPO_HIDDEN_STATES_WRAPPED_ATTR",
|
|
"_UNSLOTH_GRPO_HIDDEN_STATES_WARNING_ATTR",
|
|
"_UNSLOTH_GRPO_HIDDEN_STATES_DEGRADED_ATTR",
|
|
)
|
|
|
|
|
|
def load_rl_wrapper(names = WRAPPER_NAMES):
|
|
"""Return ``{name: object}`` for the wrapper helpers, exec'd from live source."""
|
|
text = SOURCE_PATH.read_text(encoding = "utf-8")
|
|
tree = ast.parse(text, filename = str(SOURCE_PATH))
|
|
|
|
wanted = []
|
|
for node in tree.body:
|
|
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name in names:
|
|
wanted.append(node)
|
|
elif isinstance(node, ast.Assign) and any(
|
|
isinstance(t, ast.Name) and t.id in CONSTANT_NAMES for t in node.targets
|
|
):
|
|
wanted.append(node)
|
|
|
|
found = {
|
|
node.name for node in wanted if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
|
|
}
|
|
missing = set(names) - found - set(OPTIONAL_NAMES)
|
|
names = tuple(name for name in names if name in found)
|
|
if missing:
|
|
raise AssertionError(f"missing module-level defs in {SOURCE_PATH}: {sorted(missing)}")
|
|
|
|
namespace: dict = {
|
|
"os": os,
|
|
"collections": collections,
|
|
"inspect": inspect,
|
|
"logger": logging.getLogger("unsloth-repro"),
|
|
}
|
|
exec(compile(ast.Module(body = wanted, type_ignores = []), str(SOURCE_PATH), "exec"), namespace)
|
|
return {name: namespace[name] for name in names}
|