vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-10-28 17:52:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 38

Token budget: weight split-chunk TRL metrics by token count (#13978)

Browse source 
## Summary
With `UNSLOTH_TOKEN_BUDGET` set, a micro-batch over the budget runs
`training_step` once per chunk. TRL's SFTTrainer appends one
`mean_token_accuracy` and one `entropy` entry per `compute_loss` call,
then logs the plain mean of those entries. Unequal chunks were therefore
weighted equally. In one example the logged accuracy was 0.5, while the
same unsplit batch logged 0.25. Loss and gradients were already exact.

This change folds the per-chunk entries back into one entry per
micro-batch, weighted by the token count TRL divides by:
- accuracy: shifted labels that are not -100 (TRL 0.22.2 and 1.x)
- entropy: the label mask on fused-head TRL 1.x, and the attention mask
(plus prompt-tuning virtual tokens) or every position when padding-free
on TRL 0.22.2
- DDP: one all_reduce of the counts, only on split steps, matching TRL's
gather

## Scope
- Only the split branch changes. When the budget is unset, a step is not
split, or the trainer has no `_metrics`, the code path is the same as
before.
- Unsloth's generated SFT trainer does not append these metrics, so the
default Unsloth SFT path is unaffected. The fix covers unrewritten TRL
trainers (for example `UNSLOTH_ALLOW_CPU=1`, compile fallback) and
subclasses that log TRL-style metrics.

## Tests
- `tests/test_token_budget_metrics.py`: unequal chunks against TRL
0.22.2-style and 1.x-style metric code. 2 fail before this change and
all 4 pass after.
- Existing token budget, formatter and encoding tests: 100 passed, 2
skipped.
- Real TRL 1.15 SFTTrainer: logged accuracy for a split step is 0.25 (it
was 0.5), equal to the unsplit step. Entropy also matches.
This commit is contained in:
Daniel Han 2026-10-16 05:18:22 -07:00 • committed by GitHub
parent d4e005e73e
commit f046663708
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 333 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

223
tests/test_token_budget_metrics.py Normal file
View file

@ -0,0 +1,223 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""UNSLOTH_TOKEN_BUDGET keeps TRL's per-forward ratio metrics (token accuracy, entropy) equal to an
unsplit step when a micro-batch is split into unequal chunks. CPU only."""
from __future__ import annotations
import contextlib
import importlib.util
import sys
from collections import defaultdict
from pathlib import Path
import pytest
torch = pytest.importorskip("torch")
_PATH = Path(__file__).resolve().parents[1] / "unsloth" / "utils" / "token_budget.py"
_spec = importlib.util.spec_from_file_location("_unsloth_token_budget_metrics", _PATH)
tb = importlib.util.module_from_spec(_spec)
sys.modules[_spec.name] = tb
_spec.loader.exec_module(tb)
class _Accelerator:
distributed_type = "NO"
num_processes = 1
@contextlib.contextmanager
def no_sync(self, model):
yield
def _token_stats(input_ids):
# Stand-ins for the model: a token is "predicted correctly" when even, its entropy is its id.
return (input_ids % 2 == 0).double(), input_ids.double()
class _TRL022Trainer:
"""TRL 0.22.2 SFTTrainer.compute_loss metrics: accuracy over shifted labels != -100, entropy over
the attention mask."""
num_virtual_tokens = 0
def __init__(self):
self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)}
self.accelerator = _Accelerator()
self.args = type("Args", (), {"gradient_accumulation_steps": 1})()
self.model_accepts_loss_kwargs = True
self.label_smoother = None
def compute_loss(self, inputs):
correct, entropy = _token_stats(inputs["input_ids"])
mask = inputs["attention_mask"].double()
self._metrics["train"]["entropy"].append(float((entropy * mask).sum() / mask.sum()))
self._metrics["train"]["num_tokens"] = [0]
shift_labels = inputs["labels"][..., 1:]
valid = shift_labels != -100
hits = (correct[..., 1:].bool() & valid).sum()
self._metrics["train"]["mean_token_accuracy"].append(float(hits / valid.sum()))
def training_step(
self,
model,
inputs,
num_items_in_batch = None,
):
self.compute_loss(inputs)
return torch.tensor(1.0)
class _TRL1Trainer(_TRL022Trainer):
"""TRL 1.x SFTTrainer.compute_loss metrics: both ratios over the fused head's label_mask."""
def compute_loss(self, inputs):
correct, entropy = _token_stats(inputs["input_ids"])
label_mask = inputs["labels"][..., 1:] != -100
num_valid = label_mask.sum()
entropy_sum = (entropy[..., 1:] * label_mask).sum()
num_correct = (correct[..., 1:] * label_mask).sum()
self._metrics["train"]["entropy"].append(float(entropy_sum / num_valid))
self._metrics["train"]["mean_token_accuracy"].append(float(num_correct / num_valid))
def _padded_batch():
lengths = [2, 3, 3, 12]
width = max(lengths)
gen = torch.Generator().manual_seed(0)
input_ids = torch.randint(1, 50, (len(lengths), width), generator = gen)
attention_mask = torch.zeros(len(lengths), width, dtype = torch.long)
for row, n in enumerate(lengths):
attention_mask[row, :n] = 1
labels = input_ids.masked_fill(attention_mask == 0, -100)
labels[:, 0] = -100
return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}
def _logged(trainer):
return {k: sum(v) / len(v) for k, v in trainer._metrics["train"].items()}
@pytest.mark.parametrize("trainer_cls", [_TRL022Trainer, _TRL1Trainer])
def test_split_metrics_match_unsplit_step(monkeypatch, trainer_cls):
monkeypatch.delenv("UNSLOTH_TOKEN_BALANCE", raising = False)
batch = _padded_batch()
monkeypatch.delenv("UNSLOTH_TOKEN_BUDGET", raising = False)
reference = trainer_cls()
tb.install_token_budget(reference)
reference.training_step(None, dict(batch), 10)
expected = _logged(reference)
monkeypatch.setenv("UNSLOTH_TOKEN_BUDGET", "36")
trainer = trainer_cls()
tb.install_token_budget(trainer)
chunks = tb.split_batch_by_token_budget(dict(batch), 36)
assert len(chunks) > 1 and len({c["input_ids"].shape[0] for c in chunks}) > 1
# Two steps: one entry per micro-batch, each weighted like the unsplit forward.
trainer.training_step(None, dict(batch), 10)
trainer.training_step(None, dict(batch), 10)
assert trainer._unsloth_token_budget_splits == 2
for key in ("entropy", "mean_token_accuracy"):
assert len(trainer._metrics["train"][key]) == 2
logged = _logged(trainer)
for key in ("entropy", "mean_token_accuracy"):
assert logged[key] == pytest.approx(expected[key], rel = 1e-6)
assert logged.get("num_tokens") == expected.get("num_tokens")
def test_unsplit_steps_leave_metrics_untouched(monkeypatch):
monkeypatch.delenv("UNSLOTH_TOKEN_BALANCE", raising = False)
monkeypatch.setenv("UNSLOTH_TOKEN_BUDGET", "100000")
trainer = _TRL1Trainer()
tb.install_token_budget(trainer)
trainer.training_step(None, _padded_batch(), 10)
assert trainer._unsloth_token_budget_splits == 0
assert len(trainer._metrics["train"]["entropy"]) == 1
def test_trainers_without_metrics_still_split(monkeypatch):
monkeypatch.delenv("UNSLOTH_TOKEN_BALANCE", raising = False)
monkeypatch.setenv("UNSLOTH_TOKEN_BUDGET", "36")
trainer = _TRL1Trainer()
del trainer._metrics
seen = []
trainer.training_step = lambda model, inputs, n = None: seen.append(inputs) or torch.tensor(1.0)
tb.install_token_budget(trainer)
loss = trainer.training_step(None, _padded_batch(), 10)
assert len(seen) > 1 and float(loss) == len(seen)
class _DelegatingTRL1Trainer(_TRL1Trainer):
"""A user subclass of a TRL 1.x trainer whose compute_loss only delegates to super()."""
def compute_loss(self, inputs):
return super().compute_loss(inputs)
def test_delegating_subclass_uses_parent_metric_code(monkeypatch):
monkeypatch.delenv("UNSLOTH_TOKEN_BALANCE", raising = False)
batch = _padded_batch()
monkeypatch.delenv("UNSLOTH_TOKEN_BUDGET", raising = False)
reference = _DelegatingTRL1Trainer()
tb.install_token_budget(reference)
reference.training_step(None, dict(batch), 10)
expected = _logged(reference)
monkeypatch.setenv("UNSLOTH_TOKEN_BUDGET", "36")
trainer = _DelegatingTRL1Trainer()
tb.install_token_budget(trainer)
trainer.training_step(None, dict(batch), 10)
assert trainer._unsloth_token_budget_splits == 1
logged = _logged(trainer)
for key in ("entropy", "mean_token_accuracy"):
assert len(trainer._metrics["train"][key]) == 1
assert logged[key] == pytest.approx(expected[key], rel = 1e-6)
def _per_chunk_entries(trainer_cls, batch, **attrs):
# What TRL logs with no folding: one entry per chunk.
trainer = trainer_cls()
for name, value in attrs.items():
setattr(trainer, name, value)
for chunk in tb.split_batch_by_token_budget(dict(batch), 36):
trainer.compute_loss(chunk)
return {k: list(trainer._metrics["train"][k]) for k in ("entropy", "mean_token_accuracy")}
def _split_once(monkeypatch, trainer_cls, batch, **attrs):
monkeypatch.delenv("UNSLOTH_TOKEN_BALANCE", raising = False)
monkeypatch.setenv("UNSLOTH_TOKEN_BUDGET", "36")
trainer = trainer_cls()
for name, value in attrs.items():
setattr(trainer, name, value)
tb.install_token_budget(trainer)
trainer.training_step(None, dict(batch), 10)
assert trainer._unsloth_token_budget_splits == 1
return {k: list(trainer._metrics["train"][k]) for k in ("entropy", "mean_token_accuracy")}
def test_prompt_learning_leaves_entries_unfolded(monkeypatch):
batch = _padded_batch()
expected = _per_chunk_entries(_TRL022Trainer, batch, num_virtual_tokens = 2)
assert len(expected["entropy"]) > 1
assert _split_once(monkeypatch, _TRL022Trainer, batch, num_virtual_tokens = 2) == expected
def test_ambiguous_label_source_leaves_entries_unfolded(monkeypatch):
batch = _padded_batch()
batch["shift_labels"] = batch["labels"][..., 1:].contiguous()
batch["shift_labels"] = torch.nn.functional.pad(batch["shift_labels"], (0, 1), value = -100)
expected = _per_chunk_entries(_TRL1Trainer, batch)
assert len(expected["entropy"]) > 1
assert _split_once(monkeypatch, _TRL1Trainer, batch) == expected
def test_multi_process_leaves_entries_unfolded(monkeypatch):
batch = _padded_batch()
expected = _per_chunk_entries(_TRL022Trainer, batch)
accelerator = _Accelerator()
accelerator.num_processes = 2
assert _split_once(monkeypatch, _TRL022Trainer, batch, accelerator = accelerator) == expected

110
unsloth/utils/token_budget.py
View file

@ -490,6 +490,109 @@ def _same_count_on_every_rank(count, accelerator):
return int(bounds[0]) == -int(bounds[1])
def _label_count(chunk):
"""Tokens TRL scores for token accuracy: shifted labels that are not -100."""
labels = chunk.get("shift_labels")
if labels is None:
labels = chunk.get("labels")
if labels is None:
return 0
labels = labels[..., 1:]
return int((labels != -100).sum())
def _metric_style(trainer):
"""How the trainer's TRL compute_loss averages entropy: "labels" (fused-head TRL, label mask) or
"attention" (older TRL, attention mask). None when no class in the MRO computes TRL's metrics."""
style = getattr(trainer, "_unsloth_token_budget_metric_style", False)
if style is not False:
return style
style = None
# Subclasses that only delegate to super() do not hold the metric code: use the class that does.
for cls in type(trainer).__mro__:
function = cls.__dict__.get("compute_loss")
if function is None:
continue
try:
source = inspect.getsource(function)
except Exception:
break
if "mean_token_accuracy" in source:
style = "labels" if "label_mask" in source else "attention"
break
trainer._unsloth_token_budget_metric_style = style
return style
def _entropy_count(trainer, chunk):
"""Tokens TRL averages entropy over: the label mask on fused-head TRL, else the attention mask or
every position when padding-free."""
if _metric_style(trainer) == "labels":
return _label_count(chunk)
if "attention_mask" in chunk:
return int(chunk["attention_mask"].sum())
if "position_ids" in chunk:
return int(chunk["position_ids"].numel())
return 0
# Per-forward ratio metrics TRL appends once per compute_loss call and later averages per entry.
_RATIO_METRICS = {
"mean_token_accuracy": lambda trainer, chunk: _label_count(chunk),
"entropy": _entropy_count,
}
def _is_prompt_learning(trainer, model):
if int(getattr(trainer, "num_virtual_tokens", 0) or 0) > 0:
return True
for candidate in (model, getattr(trainer, "model", None)):
config = getattr(candidate, "active_peft_config", None)
if getattr(config, "is_prompt_learning", False):
return True
return False
def _can_merge_metrics(trainer, model, chunks, accelerator):
"""Fold per-chunk entries only when token weighting provably reproduces TRL's unsplit value:
one process (TRL's cross-rank reductions differ per version), no prompt-learning PEFT (virtual
token logits are counted differently per version), one label source, and known TRL metric code."""
if accelerator is not None and int(getattr(accelerator, "num_processes", 1) or 1) > 1:
return False
if _is_prompt_learning(trainer, model):
return False
if any("shift_labels" in chunk and "labels" in chunk for chunk in chunks):
return False
return _metric_style(trainer) is not None
def _metric_lengths(trainer):
metrics = getattr(trainer, "_metrics", None)
metrics = metrics.get("train") if isinstance(metrics, Mapping) else None
if not isinstance(metrics, Mapping):
return None, None
lengths = {
key: len(metrics[key])
for key in _RATIO_METRICS
if key in metrics and isinstance(metrics[key], list)
}
return metrics, lengths
def _merge_chunk_metrics(trainer, metrics, lengths, chunks):
"""Fold the per-chunk ratio entries back into one entry for the micro-batch, weighting each chunk by
the tokens TRL divided by, so the logged value matches an unsplit step."""
keys = [key for key in _RATIO_METRICS if key in metrics and isinstance(metrics[key], list)]
keys = [key for key in keys if len(metrics[key]) - lengths.get(key, 0) == len(chunks)]
for key in keys:
row = [_RATIO_METRICS[key](trainer, chunk) for chunk in chunks]
start = lengths.get(key, 0)
values = metrics[key][start:]
total = sum(row)
merged = sum(v * w for v, w in zip(values, row)) / total if total > 0 else 0.0
metrics[key][start:] = [merged]
def install_token_budget(trainer):
"""Install whichever of UNSLOTH_TOKEN_BALANCE / UNSLOTH_TOKEN_BUDGET is set on `trainer`."""
balance = os.environ.get(_TOKEN_BALANCE_ENV, "0").lower() in ("1", "true", "yes", "on")
@ -548,6 +651,11 @@ def install_token_budget(trainer):
if len(chunks) == 1:
return original_training_step(model, inputs, num_items_in_batch, *args, **kwargs)
trainer._unsloth_token_budget_splits += 1
metrics, lengths = (
_metric_lengths(trainer)
if _can_merge_metrics(trainer, model, chunks, accelerator)
else (None, None)
)
total = None
for i, chunk in enumerate(chunks):
context = (
@ -558,6 +666,8 @@ def install_token_budget(trainer):
with context:
loss = original_training_step(model, chunk, num_items_in_batch, *args, **kwargs)
total = loss if total is None else total + loss
if metrics is not None:
_merge_chunk_metrics(trainer, metrics, lengths, chunks)
return total
trainer.training_step = training_step