mirror of
https://github.com/unslothai/unsloth.git
synced 2026-08-19 13:53:58 +00:00
7b13a787cf
* 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>
|
||
|---|---|---|
| .. | ||
| __init__.py | ||
| aime_eval.md | ||
| aime_eval.py | ||
| cleanup_utils.py | ||
| data_utils.py | ||
| generate_dataset_with_none.py | ||
| hf_utils.py | ||
| ocr_eval.md | ||
| ocr_eval.py | ||
| os_utils.py | ||
| perplexity_eval.md | ||
| perplexity_eval.py | ||
| run_none_detect_tests.py | ||
| test_attention_dispatch_dora_dtype.py | ||
| test_attention_masks.py | ||
| test_attn_mask_compat.py | ||
| test_attn_mask_upstream_drift.py | ||
| test_batched_leftpad_generation_gpu.py | ||
| test_dataset_num_proc.py | ||
| test_packing.py | ||
| test_prepare_inputs_leftpad.py | ||
| test_q_galore.py | ||
| test_qat.py | ||
| test_rope_scaling_drift.py | ||
| test_train_on_responses_only_num_proc.py | ||
| test_trunc_normal_patch.py | ||
| test_truncation_attestation.py | ||
| test_varlen_int32_overflow_guard.py | ||
| test_xformers_capability_gate.py | ||