[staging CI] unslothai/unsloth#8890 - #318
Closed
danielhanchen wants to merge 11 commits into
Closed
Conversation
TRL prepares the whole train_dataset in the SFTTrainer constructor and never consults max_steps, so a 30-step run over a large corpus tokenizes millions of rows to read a few hundred. On unsloth/open_math_reasoning a 30-step run of Qwen3-0.6B spent 11m14s in preprocessing against 1m54s of training. The reachable row count is known before any of that: steps x batch x accumulation. Bound the dataset to it, with slack, before the formatting and tokenization passes. Shuffled rather than head-sliced, so a corpus ordered by source or difficulty does not become one homogeneous slab. Skipped for epoch-bounded runs, streaming datasets, an explicit train-split range, and packing, where rows per step is unknown.
for more information, see https://pre-commit.ci
…cking
The bound was inert on MLX. _MLXTrainerAdapter stashed max_train_rows into its
dataset config, but _build_training_worker_config is a key whitelist and dropped
both keys, and _run_mlx_training loads its own dataset and applied only the
explicit slice. An Apple Silicon max_steps run still formatted the whole corpus.
The opt-out also read the requested packing value rather than the effective one.
The image, audio-codec and audio-VLM branches train without packing whatever the
config says, and the frontend hides the packing control for image VLMs without
resetting it, so a stale flag cost those runs the optimization for nothing.
The helper moves to core/training/dataset_bounds.py, which imports no torch:
the MLX worker runs on hosts that need no torch stack, so it cannot reach this
through core.training.trainer, which imports torch and unsloth at module scope.
trainer.py re-exports the names it exported before. The MLX worker recomputes
the bound from the config rather than receiving a copy, so there is one source
of truth for it, and the adapter no longer forwards a value that was dropped.
Also:
- Coerce the helper's inputs. max(1, batch_size) raised on a None or a string,
which the request schema rules out but the DB, resumed-run records and direct
callers do not, and float("inf") escaped int() as OverflowError. A row bound
is an optimization; it must never be the thing that raises.
- 0 is a legitimate seed, so seed coercion rejects only non-integers and
negatives, which numpy refuses.
- Guard the apply site on shuffle/select rather than on len(): a DatasetDict
answers len() with its split count, and an IterableDataset has no len() at all.
- Skip the bound when resuming a checkpoint that trained on the full dataset.
Trainer fast-forwards by batch count over the current dataloader
(ignore_data_skip defaults to False), so bounding a pre-bound checkpoint now
would continue it into unrelated rows. trainer_state.json records global_step
and a fractional epoch, which recovers the row count it trained on.
Tests cover the effective-packing matrix, the coercions, seed determinism and
seed 0, the DatasetDict and streaming guards, the exact-size boundary, the eval
carve leaving enough rows for max_steps, the resume detection, and the wiring of
both loaders, which no GPU-less or Apple-less CI run can otherwise reach.
for more information, see https://pre-commit.ci
Both trainers resume by jumping to a batch INDEX, not by remembering which rows they saw: HF Trainer replays the current dataloader (ignore_data_skip defaults to False), and unsloth_zoo's MLXTrainer resolves a cursor through a schedule rebuilt from whatever dataset it is handed, with no dataset-identity check on either side. So the subset a run trains on is training state, and it has to be fixed at the first start rather than derived again later from a config the user can edit between runs. The previous commit inferred it from trainer_state.json. That reads the row count exactly, but it reads the wrong number in three ways: it recovers rows rounded up to a multiple of the batch size, a partial last accumulation cycle inflates the step count, and under DDP train_batch_size omits the world size, which underestimates and so fails to fire in the unsafe direction. It also only fired above twice the bound, so a legacy checkpoint over a dataset between one and two times the bound was misread as already bounded and resumed onto a different subset. A marker written beside the checkpoints replaces the arithmetic with a recorded fact: no marker means the checkpoint predates the bound and the run continues unbounded, exactly as it did before this feature existed. The MLX loader reads and writes it too, which it did not do for the inference version. Also drop the re-export block from trainer.py: the repo's import-hoist check rejects an added-but-unused import outside a package __init__, so the tests take the helpers from the module that owns them.
for more information, see https://pre-commit.ci
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
danielhanchen
force-pushed
the
pr-8890-xplat-ci
branch
from
August 15, 2026 14:54
75167b3 to
dcf8187
Compare
…w-text packing Three follow-ups on the marker. A resume rewrites a marker that is already valid, and it did so in place. A truncating open that then fails -- a full disk is the ordinary way, and the write is best-effort so the error is swallowed -- left an empty file, which reads as "no marker" and resumes the run over the whole dataset. It is written through a temporary file and moved into place now; os.replace is atomic on POSIX and on Windows. The run directory was found by stripping any basename starting with "checkpoint-". Trainer writes checkpoint-<global_step> and nothing else under that prefix, so a run directory whose own name starts with it (a model called checkpoint-something reaches the default run name) had its marker filed one level above where the resume then looked. Only the exact shape counts. Effective packing keyed on the dataset flags alone. Raw-text and CPT runs take the text path however the dataset is flagged, since the vision and audio-VLM branch is gated on `not raw_text_mode`, and that path honours the requested value: an image or audio dataset trained raw with packing on really does pack, so it keeps the opt-out.
danielhanchen
force-pushed
the
pr-8890-xplat-ci
branch
from
August 15, 2026 15:08
f6e305f to
06d1a5e
Compare
for more information, see https://pre-commit.ci
danielhanchen
force-pushed
the
pr-8890-xplat-ci
branch
from
August 15, 2026 15:10
8d07e12 to
06d1a5e
Compare
is_dataset_image and is_dataset_audio are client-supplied and true on a column NAME match: the trainer says so itself, and keeps _dataset_has_audio_column as the tiebreaker precisely because the flag lies. A text model with a column called "audio" carries the flag, trains on the text path, and that path honours packing, so exempting it from the opt-out bounded a run that really does pack. Raw-text and CPT reach the same path from the other direction. Three rounds, three different ways for the same guess to be wrong. Only an explicit is_vlm establishes the branch now, because it means the caller probed the model and the dataset and landed on the vision branch, which sets no packing at all. The MLX loader has that; the CUDA worker does not, and the honest consequence is that a requested packing keeps its dataset unbounded there, as it did before this feature. is_vision_model spawns a subprocess and reads configs, so there is no cheap way to learn the branch at that point, and a wrong guess costs rows the run actually needed. record_row_bound now reports whether it wrote. A marker that cannot be written leaves a run whose later resume reads it as unbounded; the callers log that rather than failing a training run over it, and there is nothing to fall back to at that point anyway, since the dataset is already bounded by the time the output directory exists.
danielhanchen
force-pushed
the
pr-8890-xplat-ci
branch
from
August 15, 2026 15:23
d5243ec to
5b4ec48
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Disposable CI run for unslothai/unsloth#8890. Do not merge; closed after CI.