vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-08-21 06:43:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 28
unsloth/scripts/online_tokenization_ab.py
Daniel Han 7b13a787cf
Studio: tokenize the dataset online for plain-text single-pass runs (#8960) 
* Studio: tokenize the dataset online for plain-text single-pass runs

TRL's tokenizing map is the largest fixed cost of starting a text run:
71s of the 78s of preparation on 100k rows of OpenMathReasoning, with
dataset_num_proc already at 8. It is per-row string work, so it can run
in DataLoader workers while the GPU is busy instead of blocking the
start.

Four parts, all needed together:

  * datasets.with_transform attaches a batched tokenizer that runs on
    __getitem__. with_transform, not set_transform: the caller's split is
    also held by the preview and the row-count checks.
  * TRL gets dataset_kwargs = {"skip_prepare_dataset": True} so it does
    not run its own map over the view. Feature-detected on SFTConfig's
    fields plus SFTTrainer.__init__'s source, never assumed.
  * dataloader_num_workers / prefetch_factor / persistent_workers, sized
    from the same shared policy that sizes dataset_num_proc and capped at
    four.
  * a prewarm barrier inside _preflight_first_batch, which already built
    a loader and pulled a batch. It now drains max(grad_accum,
    workers * prefetch) microbatches, and memoizes the train loader --
    transformers caches only the eval ones, so without that train() forks
    a second worker set and drops everything the barrier filled.

The transform reproduces unsloth_zoo's sft_prepare_dataset tokenize step
exactly: same truncation and max_length, the same double-BOS rule, and
the tokenizer's whole output rather than input_ids alone, because the
collator and the attention dispatcher both branch on which keys are
present.

Default ON only for: Linux, plain text, plain tokenizer, map-style
datasets.Dataset, packing off, no custom collator, no completion masking,
not already tokenized, no token_type_ids, a raw eval split or none, at
least 10k rows, and at most one pass over the data. Everything else takes
today's path with config_args and the dataset wrapper untouched, and any
failure in the gate or the attach degrades the same way.
UNSLOTH_STUDIO_ONLINE_TOKENIZATION=0 forces it off; =1 lifts the two cost
gates but never a correctness gate.

The one-pass rule is what the measurements support: within a single pass
the workers stay ahead and there is no steady-state cost (225.45s eager
vs 225.33s online over 200 steps), while a lazy view re-tokenizes on
every further pass where Arrow would just be read.

rl.py: a split may now attest its own truncation width through
_unsloth_truncated_to, and the max_length enforcement believes it instead
of scanning. Scanning a lazily-tokenizing split reads every row, which is
the whole eager tokenize pass again, run inside __init__ where nothing
overlaps it -- and the fallback it would then take turns padding-free
off. Both copies of the scan honour it, the module-level one and the one
inlined into every generated trainer.

Measured on one B200, Qwen3-0.6B + LoRA, 100k rows, cold datasets cache:
preparation 71.2s -> 0.4s, time to first step 91.9s -> 17.7s. Losses
match: the largest per-step gap between the eager and online arms is
7e-4, smaller than the 9e-4 between two eager runs of the same seed.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Shut the online path's workers down, and refuse the rows it would fail late on

Review of the online tokenization path found two things it got wrong once it
was running, rather than in the gate.

The persistent DataLoader workers were never torn down. Persistence is what
lets the prewarm barrier's workers survive into `train()`, and the memo holds
the loader that owns them, so after `train()` returned nothing dropped the
last reference: four `pt_data_worker` processes, 5.27 GB resident between
them, still alive through merging, quantizing and GGUF export, each a fork of
a process that had already initialised CUDA. Measured on one B200, Qwen3-0.6B
+ LoRA, 12k rows: 4 workers still there fifteen seconds after training ended,
and only the process exiting cleared them. `release_train_dataloader` shuts
them down and puts the real `get_train_dataloader` back, called from a
`finally` around `train()` so it runs before `_finalize_training` rather than
after it, and again from the outer `finally` for the paths that return before
training starts -- the preflight error has already forked the workers. Same
measurement after: 4 workers before the release, 0 after. An accelerate
wrapper and the loader inside it share one iterator, so the walk counts a
worker set once and clears the reference on both.

The second is the one asymmetry the gate did not cover. A null or non-string
row fails the eager map inside the trainer constructor, in seconds, before
anything else has happened; the lazy view reads a row only when the sampler
draws it, so the same dataset trained twenty steps and exited clean, and
would have died at whatever step drew row 137. That is the one way this
feature can make a failing run worse rather than slower. Both checks are
metadata -- the dtype off the schema, `null_count` off Arrow's per-chunk
statistics -- so neither reads a row, and a `select`ed split over-reports,
which vetoes a split that might have been fine and never the reverse. No
runtime fallback on top: switching a running job to the eager path would
tokenize the whole split mid-run and hide the bad data, where an error naming
the transform says what is actually wrong.

Also:

- The Linux gate tested `sys.platform`, but the hazard it names is `spawn`
  re-importing the entry point against a `sys.path` Studio modified in
  process. A Linux host whose start method is set to spawn or forkserver is
  the identical hazard and a platform check cannot see it. Read the start
  method instead, via `allow_none` and the method list, since resolving it
  the other way pins the context and makes a later `set_start_method()`
  raise.
- The transform truncated to the `max_seq_length` the user asked for, while
  the generated `__init__` reduces that to the model's own cap before
  deriving `max_length` from it. Read the same cap, or the two paths stop
  producing the same rows and the attestation claims a width nothing applied.
- Delete `prewarm_dataloader`. It was called from nowhere, and its docstring
  described tearing the loader down so the workers do not survive, which is
  the opposite of what the shipped barrier does on purpose.
- `scripts/online_tokenization_ab.py` defaulted `--dataset` and `--model` to
  paths under one workspace. `--dataset` is required now and the rest resolve
  without them.
- Note in the module docstring that the pass gate counts train passes only:
  an eval split is re-tokenized on every evaluation, where the eager map
  tokenized it once.

The gate was well covered and the mechanism was not. Neutering `attach`,
`online_config_args` and the memo while leaving the gate saying yes left 64
of 72 tests passing. `test_online_tokenization_runtime.py` pins the three
claims that needed a real DataLoader with real forked workers to establish:
the prewarm re-iterates from the start instead of continuing (a sequential
sampler makes it exact -- continuing the prewarmed iterator loses exactly
`prewarm * batch` rows and starts at the wrong one), the loader the barrier
filled is the one handed back afterwards, and the workers are gone once
training is over.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Probe the eval split's own BOS convention instead of reusing the train split's

TRL calls _prepare_dataset once per split, so the eager path derives
add_special_tokens from each split's own first row. The online path reused the
train split's answer for the eval view, which tokenizes eval differently from
the map it stands in for whenever the two splits disagree about a leading BOS.

Also correct the prewarm barrier's docstring. torch answers a second iter() on
a persistent-workers loader with _iterator._reset(), which restarts the sampler
at row 0 and drops what is in flight, so the drained batches are tokenized
again rather than handed to step 1. No rows are lost; what the barrier buys is
workers that are already forked and past their first tokenizer touch.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Release the memoized eval loader's workers alongside the train loader

dataloader_num_workers and dataloader_persistent_workers are
TrainingArguments settings, so an online run with evaluation on forks the
same workers for the eval loader. Transformers parks that prepared loader
in _eval_dataloaders (Trainer._get_dataloader, unchanged from 4.51.3
through 5.5.0) and torch never drops _iterator on a persistent-workers
loader once it has been iterated, so those workers outlived train() and sat
resident through the merge and export that the existing cleanup exists to
protect. Drain and drop the memo too.

* Stop the online tokenization tests depending on the runner's TRL and torch

Two CPU CI environments were red for reasons that had nothing to do with
what the tests cover. The gate tests read the installed TRL through
trl_supports_skip_prepare_dataset, and the CPU job installs no TRL, so
every refusal reported the missing hook instead of the gate under test.
Pin it in the autouse fixture, the way sys.platform is already pinned, and
cover the detector and its veto directly instead.

The wiring tests import UnslothTrainer, which imports torch, at module
scope, so a runner without torch failed collection and interrupted the
whole run rather than skipping the module. Guard it with importorskip, as
the runtime tests already do.

* Tighten online tokenization comments

* Route a Hugging Face dataset id through dataset_source in the A/B harness

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-08-16 04:55:53 -07:00

251 lines
9.2 KiB
Python
Raw Permalink Blame History

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Eager vs online preparation, measured through Studio's real training path.
Drives ``UnslothTrainer.load_model`` -> ``prepare_model_for_training`` ->
``load_and_format_dataset`` -> ``start_training``, so the integrated gating is
what gets measured. The arms differ only by ``UNSLOTH_STUDIO_ONLINE_TOKENIZATION``:
python scripts/online_tokenization_ab.py --arm eager --dataset <split> --out ab_eager.json
python scripts/online_tokenization_ab.py --arm online --dataset <split> --out ab_online.json
Same seed, rows and order, so per-step losses must match; a mismatch means the
lazy transform is not producing the rows the eager map produced.
"""
from __future__ import annotations
import argparse
import json
import os
import sys
import tempfile
import time
from pathlib import Path
REPO = Path(__file__).resolve().parents[1]
# Scratch root for the per-arm `datasets` cache; no machine-specific layout.
WORKSPACE = Path(os.environ.get("UNSLOTH_WORKSPACE") or tempfile.gettempdir())
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0")
os.environ.setdefault("UNSLOTH_DISABLE_STATISTICS", "1")
sys.path.insert(0, str(REPO / "studio" / "backend"))
sys.path.insert(0, str(REPO))
def main() -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--arm", choices = ("eager", "online"), required = True)
# No default path: it would only exist on one machine.
parser.add_argument(
"--dataset",
required = True,
help = "Parquet/JSONL split, or a Hugging Face dataset id, carrying a text column",
)
parser.add_argument("--model", default = "unsloth/Qwen3-0.6B", help = "Model id or local path")
parser.add_argument("--max-steps", type = int, default = 30)
parser.add_argument("--batch-size", type = int, default = 2)
parser.add_argument("--grad-accum", type = int, default = 4)
parser.add_argument("--max-seq-length", type = int, default = 2048)
parser.add_argument("--out", required = True)
parser.add_argument("--fresh-cache", action = "store_true", default = True)
parser.add_argument("--no-fresh-cache", dest = "fresh_cache", action = "store_false")
args = parser.parse_args()
# Fresh cache per run, else the eager arm just reads the other arm's
# tokenize map out of Arrow and measures a cache hit real users never get.
if args.fresh_cache:
cache = WORKSPACE / "unsloth_ab_cache" / f"{args.arm}_{int(time.time())}"
cache.mkdir(parents = True, exist_ok = True)
os.environ["HF_DATASETS_CACHE"] = str(cache)
# Set before anything imports the gate.
if args.arm == "eager":
os.environ["UNSLOTH_STUDIO_ONLINE_TOKENIZATION"] = "0"
else:
os.environ.pop("UNSLOTH_STUDIO_ONLINE_TOKENIZATION", None)
import unsloth # noqa: F401 - must precede transformers/trl
from transformers import TrainerCallback
from core.training.trainer import UnslothTrainer
start = time.perf_counter()
marks: dict = {}
def mark(name: str) -> None:
marks[name] = round(time.perf_counter() - start, 4)
print(f"[phase] {name} @ {marks[name]}s", flush = True)
trainer = UnslothTrainer()
if not trainer.load_model(
model_name = args.model,
max_seq_length = args.max_seq_length,
load_in_4bit = True,
):
print("model load failed", file = sys.stderr)
return 1
if not trainer.prepare_model_for_training(
use_lora = True,
lora_r = 16,
lora_alpha = 16,
lora_dropout = 0.0,
target_modules = [
"q_proj",
"k_proj",
"v_proj",
"o_proj",
"gate_proj",
"up_proj",
"down_proj",
],
use_gradient_checkpointing = "unsloth",
):
print("model prepare failed", file = sys.stderr)
return 1
mark("model_ready")
# `local_datasets` resolves its entries to files and rejects anything without a
# supported extension, so a Hub id has to go through `dataset_source` instead.
local_split = os.path.exists(args.dataset) or Path(args.dataset).suffix.lower() in (
".json",
".jsonl",
".csv",
".parquet",
)
result = trainer.load_and_format_dataset(
dataset_source = None if local_split else args.dataset,
format_type = "auto",
local_datasets = [args.dataset] if local_split else None,
)
if result is None:
print("dataset load failed", file = sys.stderr)
return 1
dataset, eval_dataset = result
mark("dataset_formatted")
class _Probe(TrainerCallback):
"""Wall clock at train() and at every step, plus the loss stream."""
def __init__(self):
self.losses: list = []
self.step_times: list = []
def on_train_begin(self, targs, state, control, **kwargs):
mark("train_begin")
def on_step_end(self, targs, state, control, **kwargs):
self.step_times.append(round(time.perf_counter() - start, 4))
if len(self.step_times) == 1:
mark("first_step_end")
def on_log(
self,
targs,
state,
control,
logs = None,
**kwargs,
):
if logs and "loss" in logs:
self.losses.append(logs["loss"])
probe = _Probe()
# The trainer only exists inside the worker thread, so attach on appearance.
original_preflight = trainer._preflight_first_batch
def _preflight_with_probe():
mark("trainer_built")
trainer.trainer.add_callback(probe)
error = original_preflight()
mark("prewarm_done")
return error
trainer._preflight_first_batch = _preflight_with_probe
started = trainer.start_training(
dataset = dataset,
eval_dataset = eval_dataset,
output_dir = f"ab_{args.arm}", # resolved under Studio's outputs root
num_epochs = 1,
max_steps = args.max_steps,
batch_size = args.batch_size,
gradient_accumulation_steps = args.grad_accum,
learning_rate = 2e-4,
weight_decay = 0.01,
random_seed = 3407,
max_seq_length = args.max_seq_length,
packing = False,
train_on_completions = False,
)
if not started:
print("training failed to start", file = sys.stderr)
return 1
while trainer.training_thread and trainer.training_thread.is_alive():
time.sleep(1)
trainer.training_thread.join()
mark("train_done")
progress = trainer.get_training_progress()
error = getattr(progress, "error", None)
decision = getattr(trainer, "_online_prewarm_batches", 0)
# What the trainer actually got configured with, read off the object.
observed = {}
sft = getattr(trainer, "trainer", None)
if sft is not None:
targs = getattr(sft, "args", None)
split = getattr(sft, "train_dataset", None)
fmt = getattr(split, "format", None)
observed = {
"dataloader_num_workers": getattr(targs, "dataloader_num_workers", None),
"dataloader_persistent_workers": getattr(targs, "dataloader_persistent_workers", None),
"dataloader_prefetch_factor": getattr(targs, "dataloader_prefetch_factor", None),
"dataset_kwargs": getattr(targs, "dataset_kwargs", None),
"remove_unused_columns": getattr(targs, "remove_unused_columns", None),
"padding_free": getattr(targs, "padding_free", None),
"packing": getattr(targs, "packing", None),
"dataset_num_proc": getattr(targs, "dataset_num_proc", None),
"train_split_format": fmt.get("type") if isinstance(fmt, dict) else None,
"train_split_columns": list(getattr(split, "column_names", None) or []),
"train_split_rows": len(split) if split is not None else None,
}
payload = {
"arm": args.arm,
"error": error,
"phases": marks,
"losses": probe.losses,
"step_times": probe.step_times,
"prewarm_batches": decision,
"observed": observed,
# Studio's chat-template render, which BOTH arms do eagerly.
"format_seconds": round(
marks.get("dataset_formatted", 0.0) - marks.get("model_ready", 0.0), 4
),
# Trainer construction: TRL's tokenizing map on the eager arm, nothing online.
"prep_seconds": round(
marks.get("trainer_built", 0.0) - marks.get("dataset_formatted", 0.0), 4
),
"time_to_first_step": marks.get("first_step_end"),
"steady_state_seconds": (
round(probe.step_times[-1] - probe.step_times[0], 4)
if len(probe.step_times) > 1
else None
),
}
if probe.losses:
payload["mean_loss"] = round(sum(probe.losses) / len(probe.losses), 6)
out = Path(args.out)
out.parent.mkdir(parents = True, exist_ok = True)
out.write_text(json.dumps(payload, indent = 2), encoding = "utf-8")
print(json.dumps({k: v for k, v in payload.items() if k != "step_times"}, indent = 2))
return 1 if error else 0
if __name__ == "__main__":
raise SystemExit(main())
Reference in a new issue View git blame Copy permalink