fix(unsloth studio): only preprocess the rows a max_steps run will actually use instead of preprocessing validated datasets by LeoBorcherding · Pull Request #8890 · unslothai/unsloth · GitHub
Skip to content

fix(unsloth studio): only preprocess the rows a max_steps run will actually use instead of preprocessing validated datasets - #8890

Merged
danielhanchen merged 16 commits into
unslothai:mainfrom
LeoBorcherding:fix/train-subset-dataset-for-max-steps
Aug 16, 2026
Merged

fix(unsloth studio): only preprocess the rows a max_steps run will actually use instead of preprocessing validated datasets#8890
danielhanchen merged 16 commits into
unslothai:mainfrom
LeoBorcherding:fix/train-subset-dataset-for-max-steps

Conversation

@LeoBorcherding

@LeoBorcherding LeoBorcherding commented Aug 15, 2026

Copy link
Copy Markdown
Collaborator

A max_steps run tokenizes the entire dataset before it takes a single step. On unsloth/open_math_reasoning (27 GB on disk) a 30-step run of unsloth/Qwen3-0.6B spent 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 against Qwen/Qwen3.5-4B spent 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's SFTTrainer prepares the whole train_dataset in its constructor — _prepare_dataset has four .map() call sites and the file contains no reference to max_steps, and the chat-template formatting ahead of it maps the full dataset too.

When max_steps is 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_only drops when the response template is missing. Running short is not an error — max_steps re-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_steps is unset or <= 0, when the dataset is streaming (already sliced lazily via skip()/take()), when dataset_slice_start/dataset_slice_end was 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_steps run 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_rows is covered for the unbounded case, the slack multiplier and the floor. load_and_format_dataset is 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.py passes 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.

LeoBorcherding and others added 3 commits August 14, 2026 12:23
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.
@LeoBorcherding LeoBorcherding changed the title only preprocess the rows a max_steps run will actually use fix(unsloth studio): only preprocess the rows a max_steps run will actually use instead of preprocessing validated datasets Aug 15, 2026
@LeoBorcherding
LeoBorcherding marked this pull request as ready for review August 15, 2026 03:54
chatgpt-codex-connector[bot]

This comment was marked as resolved.

…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.
@danielhanchen

Copy link
Copy Markdown
Member

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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))

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Detect VLM structure before randomizing the first row

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 👍 / 👎.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
danielhanchen added a commit to Datta0/unsloth-staging-3 that referenced this pull request Aug 15, 2026
danielhanchen added a commit to Datta0/unsloth-staging-3 that referenced this pull request Aug 15, 2026
chatgpt-codex-connector[bot]

This comment was marked as resolved.

…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.
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@danielhanchen

Copy link
Copy Markdown
Member

@codex review

danielhanchen added a commit to Datta0/unsloth-staging-3 that referenced this pull request Aug 15, 2026
danielhanchen added a commit to Datta0/unsloth-staging-3 that referenced this pull request Aug 15, 2026

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment on lines +186 to +187
except (OSError, UnicodeError, TypeError, ValueError):
return

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Abort bounding when the initial marker write fails

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 👍 / 👎.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +3046 to +3050
dataset = bound_dataset_rows(
dataset,
max_train_rows,
max_train_rows_seed,
on_bound = _log_bound,

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Filter invalid raw rows before selecting the bounded subset

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 👍 / 👎.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@danielhanchen

Copy link
Copy Markdown
Member

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, install.sh --local in each: BEFORE is this PR's merge base 6f443b5cc, AFTER is the current head dcf8187a2. Same box, same single GPU, same HF hub cache, and HF_DATASETS_CACHE emptied before each side so neither inherits the other's arrow and map caches.

Run on both sides: unsloth/Qwen3-0.6B, LoRA 4-bit, unsloth/OpenMathReasoning with train split cot (192,523 rows), max_steps 30, batch size 2, grad accum 4. That makes 240 rows reachable, so the bound lands on the 1024 floor.

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.

Dataset row, base vs head

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.

Train tab at 75 seconds, base vs head

Measured on the same two servers that were photographed, by polling GET /api/train/status:

base 6f443b5cc head dcf8187a2
rows loaded from the Hub 192,523 192,523
rows chat templated 192,523 1,024
rows tokenized 192,523 1,024
tokenizing status updates 43 3
seconds to the first optimizer step 197.9 65.9
HF_DATASETS_CACHE after the run 15.78 GB 2.39 GB

The control is the first row: both sides log Loaded dataset from Hugging Face: unsloth/OpenMathReasoning (192,523 rows), so both download and load exactly the same thing and only what gets preprocessed moves. On the head the backend also logs Bounded dataset to 1024 of 192523 rows for a max_steps run (seed 3407).

One honest note on the new status strings. Using 1024 of 192523 rows (max_steps run) and Formatting dataset (1,024 rows)... are each the current status message for well under a fifth of a second, and /api/train/status can only report what is current, so neither the UI's 3 second poll nor a much faster one reliably catches them. They are in the backend log rather than in these screenshots. What the UI does show for many seconds, and what carries the same claim, is the row count in the chat-template and tokenizer bars above.

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.
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
chatgpt-codex-connector[bot]

This comment was marked as resolved.

danielhanchen and others added 2 commits August 15, 2026 16:45
train[1000:2000] names rows exactly as dataset_slice_start and dataset_slice_end
do, and the trainer already reads it that way one branch up, but the bound saw
both numeric fields unset and resampled a selection the user had made. Both
loaders skip it now.

The raw-mode exception also sat in the wrong place. Audio preprocessing is
chosen before the raw-text bypass, and csm, snac and whisper train on plain
Trainers with no packing argument while bicodec and dac force it off, so an
audio branch never packs whatever the mode is. Only the vision and audio-VLM
branches give way to the text path when the run is raw or CPT. The decision
moves to the callers, which know which branch they are on; effective_packing is
now just "packing was asked for and this branch can do it".

Pin the encoding on the test's own worker.py read: tests/test_source_read_encoding.py
requires it, since the platform default is cp1252 on Windows and these files
gain non-ASCII bytes routinely.
danielhanchen added a commit to Datta0/unsloth-staging-3 that referenced this pull request Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 15, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 16, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 16, 2026
@chatgpt-codex-connector

Copy link
Copy Markdown

@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 16, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 16, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Aug 16, 2026
@danielhanchen
danielhanchen merged commit 13e83d0 into unslothai:main Aug 16, 2026
40 of 47 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants