fix(unsloth studio): only preprocess the rows a max_steps run will actually use instead of preprocessing validated datasets - #8890
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
LeoBorcherding
changed the title
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
…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.
|
Pushed a commit with the two review items plus a few things the testing turned up. Numbers first, since the description said the end to end run was outstanding. End to end, measured
Time to first step 175s to 41s, TRL prepare 135s to 7s, arrow cache down 85%, step time unchanged, 30 steps in every arm. Losses: base mean 0.650 Note the corpus here is 4.3 GB of Arrow rather than the 27 GB you measured, so the absolute saving is smaller than your 11m14s. Direction and mechanism reproduce. What the commit changesMLX. The bound was not merely unapplied there, it was inert. Effective packing. Agreed, with one thing worth recording: the direction of the risk is benign. Effective packing is only ever forced off relative to the config (the VLM branch never sets it, bicodec and dac force What the testing found beyond the reviewSimulations over
Checked and clean: nothing new reaches the DB or Two things left deliberately: a separate eval split is still unbounded, and the embedding trainer has its own loader and no bound. Both are follow-ups rather than blockers. Preflight is 66 passing. The training, dataset, mlx and audio suites are unchanged against the merge base; the diffusion failures on this machine are identical before and after and are missing optional deps. |
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3598593d48
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| return dataset | ||
| if total_rows <= max_train_rows: | ||
| return dataset | ||
| bounded = dataset.shuffle(seed = _seed_int(seed, 3407)).select(range(max_train_rows)) |
There was a problem hiding this comment.
When a large image/VLM dataset contains rows with null or incomplete image, caption, or message values, this shuffle can move one of those rows to index 0 before formatting. detect_vlm_dataset_structure in utils/datasets/format_detection.py inspects only next(iter(dataset)), so a dataset whose original first row was valid can now be classified as unknown and fail conversion solely because the seeded subset starts with an incomplete row; detect the structure before shuffling or make the detector inspect multiple rows.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The premise is right, detect_vlm_dataset_structure reads only next(iter(dataset)) (format_detection.py:642), but this is not a new failure class. A dataset whose original row 0 is incomplete already fails the same way today; the bound changes which datasets draw the short straw, not what happens when one does, and on uniformly formatted rows the shuffle cannot change the verdict. The durable fix is sampling several rows the way detect_dataset_format already does, which is a pre-existing robustness gap rather than this PR's.
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
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
…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.
|
@codex review |
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 06d1a5e4cc
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| except (OSError, UnicodeError, TypeError, ValueError): | ||
| return |
There was a problem hiding this comment.
The final atomic replace protects an existing marker, but this catch still silently ignores a first-run write failure after the dataset has already been reduced to a shuffled subset. If the failure is transient—for example, temporary ENOSPC that clears before a later checkpoint—the run can successfully checkpoint without a marker; row_bound_for_resume then treats it as legacy and resumes over the full naturally ordered dataset, silently continuing on unrelated rows. Either fail the run or fall back to an unbounded dataset when the initial marker cannot be persisted.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Neither remedy is available here. Falling back to an unbounded dataset is not possible at that point: the output directory is resolved after the dataset has already been loaded, formatted and bounded, so there is nothing left to fall back to. Failing the run trades a rare degraded resume for a certain lost training run, over a file that only exists to make a later resume exact. The write is atomic now, so an existing marker cannot be destroyed, and 5b4ec48 makes record_row_bound report whether it wrote so both loaders log the failure instead of swallowing it. A resume that reads such a run as unbounded is then explainable rather than silent.
| dataset = bound_dataset_rows( | ||
| dataset, | ||
| max_train_rows, | ||
| max_train_rows_seed, | ||
| on_bound = _log_bound, |
There was a problem hiding this comment.
In raw-text or CPT mode, this samples the bounded subset before prepare_raw_text_dataset filters null and non-string text values. For a large sparse dataset with valid text rows scattered among mostly invalid rows, the selected 1,024-row subset can contain too few valid rows—or none—even though the full dataset has enough usable examples; the former repeatedly cycles a tiny set for max_steps, while the latter raises the existing “all rows were null or non-string” error and prevents a run that previously succeeded. Apply the inexpensive raw-text validity filter before shuffling and selecting the bound.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The remedy costs more than the problem. Filtering before the bound is a pass over every row of the corpus, which is the cost this change exists to avoid; on the datasets it targets that is millions of rows scanned to save a sample. And the failure needs a dataset that is almost entirely null or non-string text: for 1,024 seeded rows to contain none, well over 99 percent of the corpus has to be invalid, which is a broken dataset rather than a sparse one. A user who does have one can name the rows with an explicit train-split range, which disables the bound.
|
Before/after evidence for what this changes on screen. The PR touches no frontend file, but every count the training-start overlay prints is a count of the rows the worker is preprocessing, so the overlay is where it shows. Two isolated Studio installs, Run on both sides: The Dataset row of the preparation overlay, at the last frame it is on screen. Same row, same phase, same bar. Only the denominator moves, and the denominator is what costs the two minutes.
The same page 75 seconds after the start request. The base is still behind the overlay, tokenizing rows this run can never reach ("5% (10,000/192,523), waiting for first step (0)"). The head is at step 12 of 30, 40 percent complete, 24 seconds elapsed.
Measured on the same two servers that were photographed, by polling
The control is the first row: both sides log One honest note on the new status strings. The timings are one run per side, and the head side's startup varies by a few tens of seconds between a cold and a warm home. This pair happens to be the arrangement that flatters the PR least: the AFTER install was rebuilt immediately before its run and was cold, while BEFORE was on its second run, and the gap is still about 3x. |
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.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
Step 4a probes the model and sets is_vlm, is_audio_vlm and _audio_type on the
trainer, and 4b loads the dataset, so the branch is known before the bound is
needed; the computation just sat above both and took the default. It moves below
the probe and passes what was detected, which is what the packing opt-out wanted
all along: a vision, audio-VLM or audio-codec run cannot pack, whatever the
config says, and a CUDA image VLM with a stale packing flag gets its bound back.
Renamed to branch_never_packs, since is_vlm was never the question.
Also resolve a bare relative checkpoint path. os.path.split("checkpoint-30")
returns an empty head, so the exact-match guard added with the previous commit
rejected a real checkpoint and looked for the marker inside it, which reads a
bounded run as legacy. Its run directory is the working directory.
Verified on a GPU that the move did not quietly stop the bound: 1024 of 192523
rows, 30 steps, 35s to the first step, and the same loss trajectory as before,
so the subset is unchanged.
|
@codex review |
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9dfff03d97
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if ( | ||
| (not dataset_streaming) | ||
| and dataset_slice_start is None | ||
| and dataset_slice_end is None | ||
| ): |
There was a problem hiding this comment.
When train_split is a supported explicit range such as train[1_000:2_000], both numeric slice fields remain None, so this condition applies the shuffled max-step bound even though the user already named the intended rows. This contradicts the adjacent explicit-range opt-out and can reduce a larger requested split to 1,024 randomized rows; treat bracketed split instructions like dataset_slice_start/dataset_slice_end in both CUDA and MLX loaders.
Useful? React with 👍 / 👎.
| raw_text_mode = ( | ||
| config.get("training_type") == "Continued Pretraining" or config.get("format_type") == "raw" | ||
| ) | ||
| return bool(raw_text_mode or not branch_never_packs) |
There was a problem hiding this comment.
When a real CUDA audio model is run with raw/CPT settings and a retained packing=True, raw_text_mode makes this report effective packing even though audio preprocessing is selected before the raw-text bypass: CSM, SNAC, and Whisper use plain trainers with no packing, while BiCodec and DAC explicitly force packing off. Such max-step runs therefore preprocess or encode the entire audio dataset, defeating this optimization on its most expensive paths; the raw-mode exception should apply only to branches that actually reach the text trainer.
Useful? React with 👍 / 👎.
2 participants
A
max_stepsrun tokenizes the entire dataset before it takes a single step. Onunsloth/open_math_reasoning(27 GB on disk) a 30-step run ofunsloth/Qwen3-0.6Bspent 11m14s in preprocessing against 1m54s of training, the trainer subprocess started at 16:01:18 and step 1 landed at 16:12:32. The same dataset againstQwen/Qwen3.5-4Bspent 16m16s. Those 30 steps consumed a few hundred examples; the run paid to tokenize millions.The time goes to
datasets.map(), in passes that wrote 7.8 GB of Arrow cache. TRL'sSFTTrainerprepares the wholetrain_datasetin its constructor —_prepare_datasethas four.map()call sites and the file contains no reference tomax_steps, and the chat-template formatting ahead of it maps the full dataset too.When
max_stepsis set, the rows a run can reach are known before any of that work happens:max_steps × per_device_train_batch_size × gradient_accumulation_steps. This bounds the dataset to that count, with slack, before formatting or tokenization sees it.The subset is taken as
shuffle(seed).select(...), not a head slice, a dataset ordered by difficulty or source would otherwise turn a short run into training on one homogeneous slab. Shuffling an Arrow dataset builds an indices mapping rather than rewriting data.The slack multiplier covers rows consumed without producing a step: the eval split carved off the train set, and rows
train_on_responses_onlydrops when the response template is missing. Running short is not an error —max_stepsre-reads the subset, but it would train on the same rows twice, so the bound is deliberately loose. A floor keeps small runs from subsetting to a statistically useless handful.The bound is skipped when
max_stepsis unset or<= 0, when the dataset is streaming (already sliced lazily viaskip()/take()), whendataset_slice_start/dataset_slice_endwas given (the user named the rows), and when packing is on (one packed sample spans an unknown number of rows).TRL's own escape hatch,
dataset_kwargs={"skip_prepare_dataset": True}, doesn't fit here: it requires an already-tokenized dataset and a custom collator. Studio wants TRL to do the tokenizing, just not over rows no step will reach.Note this changes which examples a
max_stepsrun sees, previously the dataloader's ordering over the full dataset, now a seeded subset, so such runs won't reproduce step-for-step against earlier versions.Testing
Unit only so far.
max_steps_dataset_rowsis covered for the unbounded case, the slack multiplier and the floor.load_and_format_datasetis covered for subsetting a large dataset, leaving a small one untouched, deferring to an explicit train-split range, and staying off when no bound is passed. Ruff clean;test_training_preflight.pypasses 54/54.The end-to-end run has not been done yet, which is why this is a draft. The timings above are measured from a run on
main, not from this branch, what's outstanding is confirming the branch actually collapses that 11m14s and that the loss curve stays in range.