From 84fdd991db28be1f26814fb399f09d293a02aabb Mon Sep 17 00:00:00 2001 From: LeoBorcherding Date: Fri, 14 Aug 2026 12:23:15 -0500 Subject: [PATCH 01/14] only preprocess the rows a max_steps run will actually use 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. --- studio/backend/core/training/trainer.py | 60 +++++++++ studio/backend/core/training/training.py | 7 + studio/backend/core/training/worker.py | 17 ++- .../backend/tests/test_training_preflight.py | 124 +++++++++++++++++- 4 files changed, 206 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 7f4d9471c..4d73f3898 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -94,6 +94,33 @@ # Streaming eval has no __len__, so an unbounded eval would iterate the whole # source on every eval step. Cap it so each evaluation terminates. STREAMING_EVAL_MAX_SAMPLES = 500 +# Slack on the row bound below. Rows are consumed by things that never produce a +# step: the eval split carved off the train set, and the rows +# train_on_responses_only drops when the response template is missing. Running +# short is not an error -- max_steps just re-reads the subset -- but it trains on +# the same rows twice, so the bound is deliberately loose. Even 4x is three +# orders of magnitude under the datasets this exists for. +MAX_STEPS_ROW_SLACK = 4 +# Below this a subset is small enough to skew a run for no meaningful saving. +MIN_MAX_STEPS_ROWS = 1024 + + +def max_steps_dataset_rows( + max_steps: int, + batch_size: int, + gradient_accumulation_steps: int, +) -> Optional[int]: + """Rows a max_steps run can reach, or None when it is unbounded. + + TRL prepares the whole train_dataset in the SFTTrainer constructor and never + looks at max_steps, so a 30-step run over a large corpus tokenizes millions + of rows to read a few hundred. The count is known before any of that: a step + draws batch_size * gradient_accumulation_steps rows. + """ + if not max_steps or max_steps <= 0: + return None + per_step = max(1, batch_size) * max(1, gradient_accumulation_steps) + return max(MIN_MAX_STEPS_ROWS, max_steps * per_step * MAX_STEPS_ROW_SLACK) def _build_report_targets(training_args) -> list[str] | str: @@ -2572,6 +2599,8 @@ def load_and_format_dataset( dataset_local_path: Optional[str] = None, dataset_revision: Optional[str] = None, require_exact_resume_resources: bool = False, + max_train_rows: Optional[int] = None, + max_train_rows_seed: int = 3407, ) -> Optional[tuple]: """ Load and prepare a dataset for training. @@ -2579,6 +2608,9 @@ def load_and_format_dataset( Strategy: format first, then split — ensures both train and eval portions are formatted and templated. + max_train_rows bounds the rows kept before formatting, for a max_steps + run that cannot reach the whole dataset; see max_steps_dataset_rows. + Returns (dataset_info, eval_dataset) or None on error; eval_dataset may be None if no eval split is available. """ @@ -3017,6 +3049,34 @@ def _apply_raw_text_prep(ds: Dataset, split_name: str) -> Dataset: status_message = f"Sliced dataset to {len(dataset)} rows (indices {start}-{end})" ) + # Before the formatting, template and tokenization passes, all of which map + # over every row: that is the cost this avoids. Skipped when the user named + # an explicit range, which is already the rows they asked for, and when + # streaming, which was bounded lazily above. + if ( + (not dataset_streaming) + and max_train_rows is not None + and dataset_slice_start is None + and dataset_slice_end is None + ): + total_rows = len(dataset) + if total_rows > max_train_rows: + # Shuffled, not the head. A corpus ordered by source or difficulty + # would otherwise make a short run train on one homogeneous slab. + # shuffle() builds an indices mapping; it does not rewrite the table. + dataset = dataset.shuffle(seed = max_train_rows_seed).select( + range(max_train_rows) + ) + logger.info( + f"Bounded dataset to {max_train_rows} of {total_rows} rows for a " + f"max_steps run (seed {max_train_rows_seed})\n" + ) + self._update_progress( + status_message = ( + f"Using {max_train_rows} of {total_rows} rows (max_steps run)" + ) + ) + if self.should_stop: logger.info("Stopped before applying chat template\n") return None diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 8db7e6462..b62974695 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -797,7 +797,12 @@ def load_and_format_dataset( dataset_local_path: Optional[str] = None, dataset_revision: Optional[str] = None, require_exact_resume_resources: bool = False, + max_train_rows: Optional[int] = None, + max_train_rows_seed: int = 3407, ) -> Optional[tuple]: + # UnslothTrainer.__new__ hands back this adapter on an MLX host, so the + # signature has to match. The CLI does its own loading; the bound rides + # along in the config it is handed. self._dataset_config = { "hf_dataset": dataset_source or "", "local_datasets": local_datasets, @@ -816,6 +821,8 @@ def load_and_format_dataset( "dataset_snapshot_path": dataset_local_path, "dataset_revision": dataset_revision, "require_exact_dataset_resource": bool(require_exact_resume_resources), + "max_train_rows": max_train_rows, + "max_train_rows_seed": max_train_rows_seed, } self.is_cpt = bool(is_cpt) self._update_progress(status_message = "Queued MLX dataset load") diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index d4f7be7e9..df022855c 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -4017,7 +4017,7 @@ def _hip_ver_at_least(major: int, minor: int) -> bool: if backend_path not in sys.path: sys.path.insert(0, backend_path) - from core.training.trainer import UnslothTrainer + from core.training.trainer import UnslothTrainer, max_steps_dataset_rows from utils.paths import ( ensure_dir, resolve_output_dir, @@ -4084,6 +4084,19 @@ def _apply_stop(save: bool) -> None: training_type = config.get("training_type", "LoRA/QLoRA") is_cpt_for_dataset = training_type == "Continued Pretraining" + # Packing opts out: one packed sample spans an unknown number of rows, so + # steps cannot be converted to a row count. Streaming and an explicit + # train-split range opt out inside load_and_format_dataset, where they live. + max_train_rows = ( + None + if config.get("packing", False) + else max_steps_dataset_rows( + config.get("max_steps", 0) or 0, + config.get("batch_size", 2), + config.get("gradient_accumulation_steps", 4), + ) + ) + def _load_training_dataset(): result = trainer.load_and_format_dataset( dataset_source = hf_dataset if hf_dataset and hf_dataset.strip() else None, @@ -4107,6 +4120,8 @@ def _load_training_dataset(): config.get("require_exact_resume_resources") or config.get("require_exact_dataset_resource") ), + max_train_rows = max_train_rows, + max_train_rows_seed = config.get("random_seed", 3407), ) if isinstance(result, tuple): loaded_dataset, loaded_eval_dataset = result diff --git a/studio/backend/tests/test_training_preflight.py b/studio/backend/tests/test_training_preflight.py index 0c99fc3cd..f2667efea 100644 --- a/studio/backend/tests/test_training_preflight.py +++ b/studio/backend/tests/test_training_preflight.py @@ -150,12 +150,20 @@ def __init__( ): self.size = size self.info = SimpleNamespace(splits = {name: object() for name in splits}) + self.shuffle_seeds = [] def __len__(self): return self.size def select(self, indices): - return _SizedDataset(len(indices), tuple(self.info.splits)) + selected = _SizedDataset(len(indices), tuple(self.info.splits)) + selected.shuffle_seeds = list(self.shuffle_seeds) + return selected + + def shuffle(self, seed = None): + shuffled = _SizedDataset(self.size, tuple(self.info.splits)) + shuffled.shuffle_seeds = [*self.shuffle_seeds, seed] + return shuffled class _SplittableDataset(_SizedDataset): @@ -462,6 +470,120 @@ def load_cached( assert cache_calls == [("train", 33), ("validation", None)] +def _cached_only_loader(monkeypatch, train, validation = None): + """A trainer whose dataset comes from cache, with remote access fatal.""" + from hub.utils import dataset_cache + + _patch_dataset_formatting(monkeypatch) + + # row_limit arrives on the explicit-slice path, which fetches end + 1 rows. + def load_cached( + repo_id, + local_path, + *, + subset, + split, + token = None, + row_limit = None, + ): + if split == "validation": + return validation + return _SizedDataset(row_limit) if row_limit else train + + def fail_remote(*args, **kwargs): + raise AssertionError("remote dataset access is not allowed") + + monkeypatch.setattr(dataset_cache, "load_cached_hf_dataset", load_cached) + monkeypatch.setattr("core.training.trainer.load_dataset", fail_remote) + monkeypatch.setattr(sys.modules["datasets"], "get_dataset_split_names", fail_remote) + return _dataset_loader_self() + + +def test_max_steps_dataset_rows_bounds_the_run(): + from core.training.trainer import ( + MAX_STEPS_ROW_SLACK, + MIN_MAX_STEPS_ROWS, + max_steps_dataset_rows, + ) + + # An epoch-bounded run reads its whole dataset, so there is nothing to bound. + assert max_steps_dataset_rows(0, 2, 4) is None + assert max_steps_dataset_rows(None, 2, 4) is None + + assert max_steps_dataset_rows(2000, 8, 16) == 2000 * 8 * 16 * MAX_STEPS_ROW_SLACK + # Small runs land on the floor rather than a statistically useless handful. + assert max_steps_dataset_rows(30, 2, 4) == MIN_MAX_STEPS_ROWS + assert max_steps_dataset_rows(1, 1, 1) == MIN_MAX_STEPS_ROWS + + +def test_max_steps_bound_subsets_before_formatting(monkeypatch): + # The whole point: 30 steps must not tokenize a corpus of 500k rows. + trainer = _cached_only_loader(monkeypatch, _SizedDataset(500_000)) + + result = trainer.load_and_format_dataset( + "org/dataset", + dataset_local_files_only = True, + dataset_local_path = "/cache/snapshot", + max_train_rows = 1024, + max_train_rows_seed = 99, + ) + + assert result is not None + bounded = result[0]["dataset"] + assert len(bounded) == 1024 + # Shuffled first: the head of a corpus ordered by source is not a sample of it. + assert bounded.shuffle_seeds == [99] + + +def test_max_steps_bound_leaves_a_small_dataset_alone(monkeypatch): + train = _SizedDataset(40) + trainer = _cached_only_loader(monkeypatch, train) + + result = trainer.load_and_format_dataset( + "org/dataset", + dataset_local_files_only = True, + dataset_local_path = "/cache/snapshot", + max_train_rows = 1024, + ) + + assert result is not None + # Untouched, so no shuffle cost and no reordering for a run that reads it all. + assert result[0]["dataset"] is train + + +def test_max_steps_bound_defers_to_an_explicit_slice(monkeypatch): + trainer = _cached_only_loader(monkeypatch, _SizedDataset(500_000)) + + result = trainer.load_and_format_dataset( + "org/dataset", + dataset_local_files_only = True, + dataset_local_path = "/cache/snapshot", + dataset_slice_start = 8, + dataset_slice_end = 32, + max_train_rows = 1024, + ) + + assert result is not None + sliced = result[0]["dataset"] + # The user named the rows; the bound must not resample them. + assert len(sliced) == 25 + assert sliced.shuffle_seeds == [] + + +def test_max_steps_bound_is_off_without_it(monkeypatch): + train = _SizedDataset(500_000) + trainer = _cached_only_loader(monkeypatch, train) + + result = trainer.load_and_format_dataset( + "org/dataset", + dataset_local_files_only = True, + dataset_local_path = "/cache/snapshot", + ) + + assert result is not None + assert result[0]["dataset"] is train + + def test_remote_train_fallback_keeps_auto_eval_remote(monkeypatch): from hub.utils import dataset_cache From b56b22d8346778f1e520c82c2b230c80da9c3bb2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 15 Aug 2026 02:10:20 +0000 Subject: [PATCH 02/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/training/trainer.py | 4 +--- studio/backend/tests/test_training_preflight.py | 6 +++++- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 4d73f3898..1d710950e 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -106,9 +106,7 @@ def max_steps_dataset_rows( - max_steps: int, - batch_size: int, - gradient_accumulation_steps: int, + max_steps: int, batch_size: int, gradient_accumulation_steps: int ) -> Optional[int]: """Rows a max_steps run can reach, or None when it is unbounded. diff --git a/studio/backend/tests/test_training_preflight.py b/studio/backend/tests/test_training_preflight.py index f2667efea..705dcdb3d 100644 --- a/studio/backend/tests/test_training_preflight.py +++ b/studio/backend/tests/test_training_preflight.py @@ -470,7 +470,11 @@ def load_cached( assert cache_calls == [("train", 33), ("validation", None)] -def _cached_only_loader(monkeypatch, train, validation = None): +def _cached_only_loader( + monkeypatch, + train, + validation = None, +): """A trainer whose dataset comes from cache, with remote access fatal.""" from hub.utils import dataset_cache From 3598593d485d75b94c40191f87f8fc5bd06298b9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 15 Aug 2026 14:04:54 +0000 Subject: [PATCH 03/14] apply the max_steps row bound on MLX too, and gate it on effective packing 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. --- .../backend/core/training/dataset_bounds.py | 178 +++++++++++++ studio/backend/core/training/trainer.py | 59 ++--- studio/backend/core/training/training.py | 8 +- studio/backend/core/training/worker.py | 56 +++- .../backend/tests/test_training_preflight.py | 250 ++++++++++++++++++ .../tests/training-start-preparation.test.ts | 1 + 6 files changed, 495 insertions(+), 57 deletions(-) create mode 100644 studio/backend/core/training/dataset_bounds.py diff --git a/studio/backend/core/training/dataset_bounds.py b/studio/backend/core/training/dataset_bounds.py new file mode 100644 index 000000000..1dc7823b1 --- /dev/null +++ b/studio/backend/core/training/dataset_bounds.py @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Row bound for a max_steps run. + +TRL prepares the whole train_dataset in the SFTTrainer constructor and never looks +at max_steps, so a 30-step run over a large corpus tokenizes millions of rows to +read a few hundred. The count is known before any of that work happens. + +This module holds no torch and no unsloth imports: both loaders use it, and the +MLX one runs on hosts where importing core.training.trainer would drag in a torch +stack that need not exist. +""" + +import json +import os +from typing import Any, Optional + +# Slack on the row bound. Rows are consumed by things that never produce a step: +# the eval split carved off the train set, and the rows train_on_responses_only +# drops when the response template is missing. Running short is not an error -- +# max_steps just re-reads the subset -- but it trains on the same rows twice, so +# the bound is deliberately loose. Even 4x is three orders of magnitude under the +# datasets this exists for. +MAX_STEPS_ROW_SLACK = 4 +# Below this a subset is small enough to skew a run for no meaningful saving. +MIN_MAX_STEPS_ROWS = 1024 + + +def _int_or(value: Any, default: int) -> int: + """Coerce a config value to an int, falling back on anything unusable. + + Studio's request schema validates these, but the worker is also driven from + the DB, from resumed-run records and by direct callers, any of which can hand + over a None or a string. A row bound is an optimization; it must never be the + thing that raises. + """ + try: + # OverflowError: json accepts Infinity without a flag, so a config column + # or a request body can carry one. + return int(value) + except (TypeError, ValueError, OverflowError): + return default + + +def _positive_int(value: Any, default: int) -> int: + """_int_or for counts, where zero and negatives are as unusable as None.""" + number = _int_or(value, default) + return number if number > 0 else default + + +def _seed_int(value: Any, default: int) -> int: + """_int_or for seeds, where 0 is legitimate but numpy rejects negatives.""" + number = _int_or(value, default) + return number if number >= 0 else default + + +def max_steps_dataset_rows( + max_steps: Any, batch_size: Any, gradient_accumulation_steps: Any +) -> Optional[int]: + """Rows a max_steps run can reach, or None when it is unbounded. + + A step draws batch_size * gradient_accumulation_steps rows. + """ + steps = _positive_int(max_steps, 0) + if steps <= 0: + return None + per_step = _positive_int(batch_size, 1) * _positive_int(gradient_accumulation_steps, 1) + return max(MIN_MAX_STEPS_ROWS, steps * per_step * MAX_STEPS_ROW_SLACK) + + +def effective_packing(config: dict, is_vlm: bool = False) -> bool: + """Whether the trainer will actually pack, not merely what was requested. + + Packing opts the bound out because one packed sample spans an unknown number + of source rows. The stored value alone overshoots: the frontend hides the + packing control for image VLMs without resetting it, API clients can submit + the combination directly, and the image, audio-codec and audio-VLM branches + all train without packing whatever the config says. + """ + if not config.get("packing", False): + return False + if is_vlm or config.get("is_dataset_image", False) or config.get("is_dataset_audio", False): + return False + return True + + +def max_train_rows_for_config(config: dict, is_vlm: bool = False) -> Optional[int]: + """The bound for a worker config, or None when the run is not bounded. + + Streaming and an explicit train-split range opt out further down, in the + loaders, where those values live. + """ + if effective_packing(config, is_vlm = is_vlm): + return None + return max_steps_dataset_rows( + config.get("max_steps", 0) or 0, + config.get("batch_size", 2), + config.get("gradient_accumulation_steps", 4), + ) + + +def checkpoint_predates_row_bound( + checkpoint_path: Any, max_train_rows: Optional[int], config: dict +) -> bool: + """Whether a checkpoint was written against a much larger dataset. + + The subset is part of training state now. Resuming a run that was started + with the same bound continues exactly, because the bound is a function of the + seed, max_steps, batch size and accumulation. A checkpoint written before the + bound existed saw the whole corpus, and Trainer fast-forwards by batch count + over the *current* dataloader (ignore_data_skip defaults to False), so + bounding it now would resume into unrelated rows. + + trainer_state.json records global_step and a fractional epoch, and epoch is + rows_seen / dataset_rows, so the dataset the checkpoint trained on can be + recovered. Anything unreadable answers False: an unresumable checkpoint is + the resume path's problem, not this one's. + """ + if not checkpoint_path or not max_train_rows: + return False + state_file = os.path.join(str(checkpoint_path), "trainer_state.json") + try: + with open(state_file, encoding = "utf-8") as handle: + state = json.load(handle) + step = float(state["global_step"]) + epoch = float(state["epoch"]) + except (OSError, UnicodeDecodeError, ValueError, TypeError, KeyError): + return False + if step <= 0 or epoch <= 0 or step != step or epoch != epoch: + return False + # The checkpoint's own batch size when it recorded one, since a resume may + # carry a different one. Accumulation is not in the state, so the current + # config answers for it; changing it on resume already makes the trainer's + # own fast-forward arithmetic unreliable, and reading it wrong here only + # drops the bound, which is the pre-existing behaviour. + per_step = _positive_int( + state.get("train_batch_size"), _positive_int(config.get("batch_size"), 1) + ) * _positive_int(config.get("gradient_accumulation_steps"), 1) + previous_rows = (step * per_step) / epoch + # Twice the bound, so an eval carve or masked-out rows in the earlier leg + # cannot read as a different dataset. + return previous_rows > max_train_rows * 2 + + +def bound_dataset_rows( + dataset, + max_train_rows: Optional[int], + seed: Any = 3407, + *, + on_bound = None, +): + """Cut a map-style dataset to max_train_rows rows, or return it untouched. + + Shuffled, not the head. A corpus ordered by source or difficulty would + otherwise make a short run train on one homogeneous slab. shuffle() builds an + indices mapping; it does not rewrite the table. + + Callers apply this before the formatting, template and tokenization passes, + all of which map over every row: that is the cost this avoids. + """ + if not max_train_rows or max_train_rows <= 0: + return dataset + # A DatasetDict answers len() with its split count, so guard on the ops this + # needs rather than on the type: anything else is left alone. + if not hasattr(dataset, "shuffle") or not hasattr(dataset, "select"): + return dataset + try: + total_rows = len(dataset) + except TypeError: + # No __len__ means a streaming dataset, which is bounded lazily instead. + return dataset + if total_rows <= max_train_rows: + return dataset + bounded = dataset.shuffle(seed = _seed_int(seed, 3407)).select(range(max_train_rows)) + if on_bound is not None: + on_bound(max_train_rows, total_rows) + return bounded diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 1d710950e..2092e93d4 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -89,37 +89,21 @@ should_use_mlx_training_backend, ) +# Re-exported: the worker and the tests import these from here, and the MLX +# loader imports them from the light module directly. +from .dataset_bounds import ( # noqa: F401 + MAX_STEPS_ROW_SLACK, + MIN_MAX_STEPS_ROWS, + bound_dataset_rows, + max_steps_dataset_rows, + max_train_rows_for_config, +) + logger = get_logger(__name__) # Streaming eval has no __len__, so an unbounded eval would iterate the whole # source on every eval step. Cap it so each evaluation terminates. STREAMING_EVAL_MAX_SAMPLES = 500 -# Slack on the row bound below. Rows are consumed by things that never produce a -# step: the eval split carved off the train set, and the rows -# train_on_responses_only drops when the response template is missing. Running -# short is not an error -- max_steps just re-reads the subset -- but it trains on -# the same rows twice, so the bound is deliberately loose. Even 4x is three -# orders of magnitude under the datasets this exists for. -MAX_STEPS_ROW_SLACK = 4 -# Below this a subset is small enough to skew a run for no meaningful saving. -MIN_MAX_STEPS_ROWS = 1024 - - -def max_steps_dataset_rows( - max_steps: int, batch_size: int, gradient_accumulation_steps: int -) -> Optional[int]: - """Rows a max_steps run can reach, or None when it is unbounded. - - TRL prepares the whole train_dataset in the SFTTrainer constructor and never - looks at max_steps, so a 30-step run over a large corpus tokenizes millions - of rows to read a few hundred. The count is known before any of that: a step - draws batch_size * gradient_accumulation_steps rows. - """ - if not max_steps or max_steps <= 0: - return None - per_step = max(1, batch_size) * max(1, gradient_accumulation_steps) - return max(MIN_MAX_STEPS_ROWS, max_steps * per_step * MAX_STEPS_ROW_SLACK) - def _build_report_targets(training_args) -> list[str] | str: report_to: list[str] = [] @@ -3053,28 +3037,25 @@ def _apply_raw_text_prep(ds: Dataset, split_name: str) -> Dataset: # streaming, which was bounded lazily above. if ( (not dataset_streaming) - and max_train_rows is not None and dataset_slice_start is None and dataset_slice_end is None ): - total_rows = len(dataset) - if total_rows > max_train_rows: - # Shuffled, not the head. A corpus ordered by source or difficulty - # would otherwise make a short run train on one homogeneous slab. - # shuffle() builds an indices mapping; it does not rewrite the table. - dataset = dataset.shuffle(seed = max_train_rows_seed).select( - range(max_train_rows) - ) + def _log_bound(kept, total): logger.info( - f"Bounded dataset to {max_train_rows} of {total_rows} rows for a " + f"Bounded dataset to {kept} of {total} rows for a " f"max_steps run (seed {max_train_rows_seed})\n" ) self._update_progress( - status_message = ( - f"Using {max_train_rows} of {total_rows} rows (max_steps run)" - ) + status_message = f"Using {kept} of {total} rows (max_steps run)" ) + dataset = bound_dataset_rows( + dataset, + max_train_rows, + max_train_rows_seed, + on_bound = _log_bound, + ) + if self.should_stop: logger.info("Stopped before applying chat template\n") return None diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index b62974695..3a62fbbcd 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -801,8 +801,10 @@ def load_and_format_dataset( max_train_rows_seed: int = 3407, ) -> Optional[tuple]: # UnslothTrainer.__new__ hands back this adapter on an MLX host, so the - # signature has to match. The CLI does its own loading; the bound rides - # along in the config it is handed. + # signature has to match. The MLX worker does its own loading and derives + # the max_steps row bound from the config it is handed, so the two bound + # arguments are accepted and deliberately not forwarded: a copy here would + # be a second source of truth that _build_training_worker_config drops. self._dataset_config = { "hf_dataset": dataset_source or "", "local_datasets": local_datasets, @@ -821,8 +823,6 @@ def load_and_format_dataset( "dataset_snapshot_path": dataset_local_path, "dataset_revision": dataset_revision, "require_exact_dataset_resource": bool(require_exact_resume_resources), - "max_train_rows": max_train_rows, - "max_train_rows_seed": max_train_rows_seed, } self.is_cpt = bool(is_cpt) self._update_progress(status_message = "Queued MLX dataset load") diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index df022855c..a080eb0cf 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -54,6 +54,13 @@ from utils.hardware import apply_gpu_ids from utils.hf_dataset_options import hf_dataset_split_instruction_names +# Light module on purpose: the MLX branch below runs on hosts that have no torch, +# so it cannot reach these through core.training.trainer. +from core.training.dataset_bounds import ( + bound_dataset_rows, + checkpoint_predates_row_bound, + max_train_rows_for_config, +) from utils.training_runs import build_default_output_dir_name from utils.wheel_utils import ( direct_wheel_url, @@ -2779,14 +2786,29 @@ def _send(event_type, **kwargs): slice_end = config.get("dataset_slice_end") config["_dataset_loaded_from_exact_snapshot"] = False + # A max_steps run cannot reach the whole dataset, and everything below this + # point -- formatting, chat templating, tokenization -- maps over every row. + # Recomputed from the config rather than carried over from the parent so an + # MLX run can never train against a bound derived from stale values. + mlx_max_train_rows = max_train_rows_for_config(config, is_vlm = is_vlm) + def _slice(ds): if slice_start is not None or slice_end is not None: start = slice_start if slice_start is not None else 0 end = slice_end if slice_end is not None else len(ds) - 1 if end < start: return ds.select([]) - ds = ds.select(range(start, min(end + 1, len(ds)))) - return ds + # The user named these rows; the bound below defers to that. + return ds.select(range(start, min(end + 1, len(ds)))) + return bound_dataset_rows( + ds, + mlx_max_train_rows, + config.get("random_seed", 3407), + on_bound = lambda kept, total: _send( + "status", + status_message = f"Using {kept} of {total} rows (max_steps run)", + ), + ) def _load_local(file_paths): from datasets import load_from_disk @@ -4017,7 +4039,7 @@ def _hip_ver_at_least(major: int, minor: int) -> bool: if backend_path not in sys.path: sys.path.insert(0, backend_path) - from core.training.trainer import UnslothTrainer, max_steps_dataset_rows + from core.training.trainer import UnslothTrainer from utils.paths import ( ensure_dir, resolve_output_dir, @@ -4084,18 +4106,24 @@ def _apply_stop(save: bool) -> None: training_type = config.get("training_type", "LoRA/QLoRA") is_cpt_for_dataset = training_type == "Continued Pretraining" - # Packing opts out: one packed sample spans an unknown number of rows, so - # steps cannot be converted to a row count. Streaming and an explicit - # train-split range opt out inside load_and_format_dataset, where they live. - max_train_rows = ( - None - if config.get("packing", False) - else max_steps_dataset_rows( - config.get("max_steps", 0) or 0, - config.get("batch_size", 2), - config.get("gradient_accumulation_steps", 4), + # Effective packing opts out: one packed sample spans an unknown number of + # rows, so steps cannot be converted to a row count. Streaming and an + # explicit train-split range opt out inside load_and_format_dataset, where + # they live. + max_train_rows = max_train_rows_for_config(config) + # A resume keeps the same rows, because the bound is a function of the + # seed, max_steps, batch size and accumulation. A checkpoint from before + # the bound existed is the exception: it trained on the whole dataset, and + # the trainer fast-forwards by batch count over the current dataloader, so + # bounding it now would continue into unrelated rows. + if checkpoint_predates_row_bound( + config.get("resume_from_checkpoint"), max_train_rows, config + ): + logger.info( + "Resuming a checkpoint that trained on the full dataset: skipping the " + "max_steps row bound so the run continues over the same rows\n" ) - ) + max_train_rows = None def _load_training_dataset(): result = trainer.load_and_format_dataset( diff --git a/studio/backend/tests/test_training_preflight.py b/studio/backend/tests/test_training_preflight.py index 705dcdb3d..d84109821 100644 --- a/studio/backend/tests/test_training_preflight.py +++ b/studio/backend/tests/test_training_preflight.py @@ -588,6 +588,256 @@ def test_max_steps_bound_is_off_without_it(monkeypatch): assert result[0]["dataset"] is train +def test_max_steps_dataset_rows_survives_unusable_numbers(): + from core.training.trainer import MIN_MAX_STEPS_ROWS, max_steps_dataset_rows + + # The worker is also driven from the DB and by direct callers, so a None or a + # string reaches this. A row bound is an optimization; it must never raise. + assert max_steps_dataset_rows(30, None, None) == MIN_MAX_STEPS_ROWS + assert max_steps_dataset_rows(30, "2", "4") == MIN_MAX_STEPS_ROWS + assert max_steps_dataset_rows("30", 2, 4) == MIN_MAX_STEPS_ROWS + assert max_steps_dataset_rows(-5, 2, 4) is None + assert max_steps_dataset_rows("not a number", 2, 4) is None + # A bound this far past any corpus is a no-op at the apply site, not an error. + assert max_steps_dataset_rows(10**9, 2, 4) == 10**9 * 8 * 4 + + +def test_effective_packing_decides_the_opt_out(): + from core.training.dataset_bounds import effective_packing, max_train_rows_for_config + + text = {"max_steps": 30, "batch_size": 2, "gradient_accumulation_steps": 4} + + # Packing spans an unknown number of rows per sample, so text runs opt out. + assert effective_packing({**text, "packing": True}) is True + assert max_train_rows_for_config({**text, "packing": True}) is None + + # ...but the image, audio and VLM branches train without packing whatever the + # config says, and the frontend hides the control without resetting it. A + # stale flag must not cost those runs the bound. + assert effective_packing({**text, "packing": True, "is_dataset_image": True}) is False + assert effective_packing({**text, "packing": True, "is_dataset_audio": True}) is False + assert effective_packing({**text, "packing": True}, is_vlm = True) is False + assert max_train_rows_for_config({**text, "packing": True}, is_vlm = True) == 1024 + assert max_train_rows_for_config({**text, "packing": True, "is_dataset_image": True}) == 1024 + + # An epoch-bounded run is unbounded whatever packing says. + assert max_train_rows_for_config({"max_steps": 0, "packing": False}) is None + + +def test_bound_dataset_rows_edges(): + from core.training.dataset_bounds import bound_dataset_rows + + class _Streaming: + """No __len__, like an IterableDataset: bounded lazily elsewhere.""" + + def shuffle(self, seed = None): + raise AssertionError("a streaming dataset must not be shuffled eagerly") + + exact = _SizedDataset(1024) + assert bound_dataset_rows(exact, 1024, 3407) is exact + assert len(bound_dataset_rows(_SizedDataset(1025), 1024, 3407)) == 1024 + + # A non-positive bound from a direct caller would select an empty dataset. + untouched = _SizedDataset(500_000) + assert bound_dataset_rows(untouched, 0, 3407) is untouched + assert bound_dataset_rows(untouched, -5, 3407) is untouched + assert bound_dataset_rows(untouched, None, 3407) is untouched + + streaming = _Streaming() + assert bound_dataset_rows(streaming, 1024, 3407) is streaming + + # A seed the config could not coerce still has to produce a subset. + assert len(bound_dataset_rows(_SizedDataset(500_000), 1024, None)) == 1024 + + +def test_bound_dataset_rows_keeps_seed_zero(): + from datasets import Dataset + + from core.training.dataset_bounds import bound_dataset_rows + + source = Dataset.from_dict({"row": list(range(5000))}) + + # 0 is a legitimate seed, not a missing one: it must not collapse onto the + # default, or every run configured with it trains on the same other subset. + assert bound_dataset_rows(source, 1024, 0)["row"] != bound_dataset_rows( + source, 1024, 3407 + )["row"] + assert bound_dataset_rows(source, 1024, 0)["row"] == bound_dataset_rows( + source, 1024, 0 + )["row"] + + +def test_bound_dataset_rows_survives_a_hostile_seed(): + from datasets import Dataset + + from core.training.dataset_bounds import bound_dataset_rows + + source = Dataset.from_dict({"row": list(range(3000))}) + + # numpy rejects a negative seed, and -1 is a common "pick one for me" + # sentinel; json accepts Infinity with no flag, so a stored config can hold + # one. Neither may take a training run down. + for seed in (-1, -3407, float("inf"), float("nan"), "3407", None, "seed"): + assert len(bound_dataset_rows(source, 1024, seed)) == 1024 + + +def test_max_steps_dataset_rows_survives_infinity(): + from core.training.dataset_bounds import MIN_MAX_STEPS_ROWS, max_steps_dataset_rows + + infinity = float("inf") + assert max_steps_dataset_rows(infinity, 2, 4) is None + assert max_steps_dataset_rows(30, infinity, 4) == MIN_MAX_STEPS_ROWS + assert max_steps_dataset_rows(30, 2, infinity) == MIN_MAX_STEPS_ROWS + + +def test_bound_dataset_rows_leaves_a_dataset_dict_alone(): + from datasets import Dataset, DatasetDict + + from core.training.dataset_bounds import bound_dataset_rows + + # len() on a DatasetDict is the split count, so the row comparison is + # meaningless there; it has no select() either. + splits = DatasetDict( + { + "train": Dataset.from_dict({"row": list(range(5000))}), + "test": Dataset.from_dict({"row": list(range(100))}), + } + ) + assert bound_dataset_rows(splits, 1024, 3407) is splits + + +def test_checkpoint_predates_row_bound(tmp_path): + import json + + from core.training.dataset_bounds import checkpoint_predates_row_bound + + config = {"batch_size": 2, "gradient_accumulation_steps": 4} + + def _checkpoint(name, *, step, epoch): + path = tmp_path / name + path.mkdir() + (path / "trainer_state.json").write_text( + json.dumps({"global_step": step, "epoch": epoch}) + ) + return str(path) + + # 15 steps x 8 rows against 192,523 rows: a run that saw the whole corpus. + full = _checkpoint("full", step = 15, epoch = 120 / 192_523) + assert checkpoint_predates_row_bound(full, 1024, config) is True + + # The same 15 steps against a 1,024-row bound: already bounded, so resuming + # keeps the bound and continues over the same rows. + bounded = _checkpoint("bounded", step = 15, epoch = 120 / 1024) + assert checkpoint_predates_row_bound(bounded, 1024, config) is False + + # The checkpoint's own batch size wins over a resume that changed it. + recorded = tmp_path / "recorded" + recorded.mkdir() + (recorded / "trainer_state.json").write_text( + json.dumps({"global_step": 15, "epoch": 120 / 192_523, "train_batch_size": 2}) + ) + assert checkpoint_predates_row_bound(str(recorded), 1024, {**config, "batch_size": 8}) is True + + # Nothing to decide from, and nothing to break: leave the bound alone. + assert checkpoint_predates_row_bound(None, 1024, config) is False + assert checkpoint_predates_row_bound(full, None, config) is False + assert checkpoint_predates_row_bound(str(tmp_path / "missing"), 1024, config) is False + empty = tmp_path / "empty" + empty.mkdir() + (empty / "trainer_state.json").write_text("{}") + assert checkpoint_predates_row_bound(str(empty), 1024, config) is False + broken = tmp_path / "broken" + broken.mkdir() + (broken / "trainer_state.json").write_text("not json") + assert checkpoint_predates_row_bound(str(broken), 1024, config) is False + + +def test_bound_dataset_rows_is_deterministic_and_seed_sensitive(): + from datasets import Dataset + + from core.training.dataset_bounds import bound_dataset_rows + + source = Dataset.from_dict({"row": list(range(5000)), "text": [f"t{i}" for i in range(5000)]}) + + first = bound_dataset_rows(source, 1024, 3407)["row"] + second = bound_dataset_rows(source, 1024, 3407)["row"] + other = bound_dataset_rows(source, 1024, 99)["row"] + + assert len(first) == 1024 + assert first == second + assert first != other + # The head of a corpus ordered by source or difficulty is not a sample of it. + assert first != list(range(1024)) + # Features survive the shuffle+select, so the formatting passes still work. + assert bound_dataset_rows(source, 1024, 3407).column_names == ["row", "text"] + + +def test_bound_leaves_enough_rows_after_the_eval_carve(): + from datasets import Dataset + + from core.training.dataset_bounds import bound_dataset_rows, max_train_rows_for_config + from core.training.eval_dataset import split_dataset_for_evaluation + + config = {"max_steps": 30, "batch_size": 2, "gradient_accumulation_steps": 4} + rows = max_train_rows_for_config(config) + source = Dataset.from_dict({"text": [f"t{i}" for i in range(500_000)]}) + + bounded = bound_dataset_rows(source, rows, 3407) + train, _eval = split_dataset_for_evaluation(bounded) + + # The eval carve is what MAX_STEPS_ROW_SLACK is budgeted for: the run must + # still reach max_steps without re-reading rows. + needed = config["max_steps"] * config["batch_size"] * config["gradient_accumulation_steps"] + assert len(train) >= needed + + +def test_both_loaders_apply_the_row_bound(): + """Guards the wiring: the helpers are useless if a loader stops calling them. + + Read from source because driving the CUDA worker needs a GPU and the MLX one + needs Apple hardware, so neither call site is otherwise reachable in CI. + """ + import ast + from pathlib import Path + + worker_src = (Path(__file__).resolve().parents[1] / "core/training/worker.py").read_text() + tree = ast.parse(worker_src) + calls = {} + for node in ast.walk(tree): + if not isinstance(node, ast.FunctionDef): + continue + names = { + sub.func.id + for sub in ast.walk(node) + if isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name) + } + calls[node.name] = names + + # The CUDA worker derives the bound and hands it to load_and_format_dataset. + assert "max_train_rows_for_config" in calls["run_training_process"] + assert "max_train_rows = max_train_rows" in worker_src + # The MLX worker loads its own dataset, so it has to bound its own rows. + assert "bound_dataset_rows" in calls["_slice"] + assert "max_train_rows_for_config" in calls["_run_mlx_training"] + + +def test_mlx_adapter_keeps_one_source_of_truth_for_the_bound(): + from core.training.training import _build_training_worker_config + + config = _build_training_worker_config( + {"model_name": "org/model", "max_steps": 30, "batch_size": 2} + ) + # The normalized worker config is a whitelist: a forwarded copy of the bound + # would be dropped here and silently disagree with what the worker computes. + assert "max_train_rows" not in config + assert "max_train_rows_seed" not in config + # Everything the worker needs to recompute it does survive. + assert config["max_steps"] == 30 + assert config["batch_size"] == 2 + assert config["gradient_accumulation_steps"] == 4 + assert config["random_seed"] == 3407 + + def test_remote_train_fallback_keeps_auto_eval_remote(monkeypatch): from hub.utils import dataset_cache diff --git a/studio/frontend/tests/training-start-preparation.test.ts b/studio/frontend/tests/training-start-preparation.test.ts index 7f2e509cc..c56402e3b 100644 --- a/studio/frontend/tests/training-start-preparation.test.ts +++ b/studio/frontend/tests/training-start-preparation.test.ts @@ -49,6 +49,7 @@ test("every status the worker sends reaches a row", () => { "Formatting VLM dataset...", "Dataset ready (1,000 samples, chatml format)", "Sliced dataset to 500 rows (indices 0-500)", + "Using 1024 of 192523 rows (max_steps run)", "Loaded 1000 samples from local files", "Encoding audio with SNAC...", 'Tokenizing ["text"] (num_proc=4) 15% (32,000/207,865)', From 80e8aa2ae0447a8b9c7f0ce4a7ad037de6e11501 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 15 Aug 2026 14:06:53 +0000 Subject: [PATCH 04/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/training/trainer.py | 2 ++ studio/backend/core/training/worker.py | 1 + studio/backend/tests/test_training_preflight.py | 14 +++++--------- 3 files changed, 8 insertions(+), 9 deletions(-) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 2092e93d4..7f353a357 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -105,6 +105,7 @@ # source on every eval step. Cap it so each evaluation terminates. STREAMING_EVAL_MAX_SAMPLES = 500 + def _build_report_targets(training_args) -> list[str] | str: report_to: list[str] = [] if training_args.get("enable_wandb", False): @@ -3040,6 +3041,7 @@ def _apply_raw_text_prep(ds: Dataset, split_name: str) -> Dataset: and dataset_slice_start is None and dataset_slice_end is None ): + def _log_bound(kept, total): logger.info( f"Bounded dataset to {kept} of {total} rows for a " diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index a080eb0cf..e93dbd467 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -54,6 +54,7 @@ from utils.hardware import apply_gpu_ids from utils.hf_dataset_options import hf_dataset_split_instruction_names + # Light module on purpose: the MLX branch below runs on hosts that have no torch, # so it cannot reach these through core.training.trainer. from core.training.dataset_bounds import ( diff --git a/studio/backend/tests/test_training_preflight.py b/studio/backend/tests/test_training_preflight.py index d84109821..19876f868 100644 --- a/studio/backend/tests/test_training_preflight.py +++ b/studio/backend/tests/test_training_preflight.py @@ -659,12 +659,10 @@ def test_bound_dataset_rows_keeps_seed_zero(): # 0 is a legitimate seed, not a missing one: it must not collapse onto the # default, or every run configured with it trains on the same other subset. - assert bound_dataset_rows(source, 1024, 0)["row"] != bound_dataset_rows( - source, 1024, 3407 - )["row"] - assert bound_dataset_rows(source, 1024, 0)["row"] == bound_dataset_rows( - source, 1024, 0 - )["row"] + assert ( + bound_dataset_rows(source, 1024, 0)["row"] != bound_dataset_rows(source, 1024, 3407)["row"] + ) + assert bound_dataset_rows(source, 1024, 0)["row"] == bound_dataset_rows(source, 1024, 0)["row"] def test_bound_dataset_rows_survives_a_hostile_seed(): @@ -716,9 +714,7 @@ def test_checkpoint_predates_row_bound(tmp_path): def _checkpoint(name, *, step, epoch): path = tmp_path / name path.mkdir() - (path / "trainer_state.json").write_text( - json.dumps({"global_step": step, "epoch": epoch}) - ) + (path / "trainer_state.json").write_text(json.dumps({"global_step": step, "epoch": epoch})) return str(path) # 15 steps x 8 rows against 192,523 rows: a run that saw the whole corpus. From 2ceb57afde7b006a870140c68470e06244019b2d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 15 Aug 2026 14:52:38 +0000 Subject: [PATCH 05/14] pin the row bound a run started with instead of inferring it on resume 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. --- .../backend/core/training/dataset_bounds.py | 118 ++++++++++++------ studio/backend/core/training/trainer.py | 10 +- studio/backend/core/training/worker.py | 44 ++++--- .../backend/tests/test_training_preflight.py | 90 +++++++------ 4 files changed, 166 insertions(+), 96 deletions(-) diff --git a/studio/backend/core/training/dataset_bounds.py b/studio/backend/core/training/dataset_bounds.py index 1dc7823b1..8e214548f 100644 --- a/studio/backend/core/training/dataset_bounds.py +++ b/studio/backend/core/training/dataset_bounds.py @@ -25,6 +25,9 @@ MAX_STEPS_ROW_SLACK = 4 # Below this a subset is small enough to skew a run for no meaningful saving. MIN_MAX_STEPS_ROWS = 1024 +# Written into a run's output directory at its first start; read back on resume. +# Its absence is the signal that a checkpoint predates the bound. +ROW_BOUND_MARKER_FILE = "unsloth_row_bound.json" def _int_or(value: Any, default: int) -> int: @@ -100,47 +103,90 @@ def max_train_rows_for_config(config: dict, is_vlm: bool = False) -> Optional[in ) -def checkpoint_predates_row_bound( - checkpoint_path: Any, max_train_rows: Optional[int], config: dict -) -> bool: - """Whether a checkpoint was written against a much larger dataset. +def run_dir_for_checkpoint(checkpoint_path: Any) -> Optional[str]: + """The run directory a checkpoint lives in, or None when there is none. - The subset is part of training state now. Resuming a run that was started - with the same bound continues exactly, because the bound is a function of the - seed, max_steps, batch size and accumulation. A checkpoint written before the - bound existed saw the whole corpus, and Trainer fast-forwards by batch count - over the *current* dataloader (ignore_data_skip defaults to False), so - bounding it now would resume into unrelated rows. + Checkpoints are written as ``/checkpoint-``; a caller that + names the run directory itself gets it back unchanged. + """ + if not checkpoint_path: + return None + path = str(checkpoint_path).rstrip("/\\") + if not path: + return None + head, tail = os.path.split(path) + if tail.startswith("checkpoint-") and head: + return head + return path + + +def record_row_bound(output_dir: Any, max_train_rows: Optional[int], seed: Any = 3407) -> None: + """Record the bound a run started with, beside its checkpoints. + + The subset a run trains on is training state: it has to be fixed at the first + start and read back on every resume, because both loaders fast-forward to a + batch *index* and the ordering is a function of the bound. Deriving it again + on resume cannot work -- the config it is derived from is editable between + runs, and a checkpoint written before this feature existed leaves no + arithmetic that distinguishes it reliably. - trainer_state.json records global_step and a fractional epoch, and epoch is - rows_seen / dataset_rows, so the dataset the checkpoint trained on can be - recovered. Anything unreadable answers False: an unresumable checkpoint is - the resume path's problem, not this one's. + Best effort by design. A run must never fail over a marker; a marker that + could not be written costs the optimization on the next resume and nothing + else. """ - if not checkpoint_path or not max_train_rows: - return False - state_file = os.path.join(str(checkpoint_path), "trainer_state.json") + run_dir = run_dir_for_checkpoint(output_dir) + if not run_dir: + return + try: + with open( + os.path.join(run_dir, ROW_BOUND_MARKER_FILE), "w", encoding = "utf-8" + ) as handle: + json.dump( + { + "max_train_rows": _positive_int(max_train_rows, 0) or None, + "seed": _seed_int(seed, 3407), + }, + handle, + ) + except (OSError, UnicodeError, TypeError, ValueError): + return + + +def row_bound_for_resume( + checkpoint_path: Any, max_train_rows: Optional[int], seed: Any = 3407 +) -> tuple[Optional[int], int]: + """The (rows, seed) a resume must use, or the freshly computed pair. + + Not resuming: the caller's own values, which record_row_bound then pins. + + Resuming a run recorded by record_row_bound: that run's values, so the rows + and their order are exactly the ones it was training on, whatever the config + now says about max_steps, batch size or accumulation. + + Resuming anything with no marker -- a checkpoint written before the bound + existed, or one whose marker is unreadable: no bound. Such a checkpoint + trained on the whole corpus in its natural order, and both trainers resume by + batch index rather than by remembering which rows they saw (HF Trainer + replays the current dataloader, `ignore_data_skip` defaulting to False; + unsloth_zoo's MLXTrainer jumps a cursor into a schedule rebuilt from the + current dataset), so a shuffled subset would silently continue on unrelated + rows. + """ + fallback_seed = _seed_int(seed, 3407) + if not checkpoint_path: + return max_train_rows, fallback_seed + run_dir = run_dir_for_checkpoint(checkpoint_path) + if not run_dir: + return max_train_rows, fallback_seed try: - with open(state_file, encoding = "utf-8") as handle: - state = json.load(handle) - step = float(state["global_step"]) - epoch = float(state["epoch"]) + with open( + os.path.join(run_dir, ROW_BOUND_MARKER_FILE), encoding = "utf-8" + ) as handle: + marker = json.load(handle) + recorded = marker["max_train_rows"] except (OSError, UnicodeDecodeError, ValueError, TypeError, KeyError): - return False - if step <= 0 or epoch <= 0 or step != step or epoch != epoch: - return False - # The checkpoint's own batch size when it recorded one, since a resume may - # carry a different one. Accumulation is not in the state, so the current - # config answers for it; changing it on resume already makes the trainer's - # own fast-forward arithmetic unreliable, and reading it wrong here only - # drops the bound, which is the pre-existing behaviour. - per_step = _positive_int( - state.get("train_batch_size"), _positive_int(config.get("batch_size"), 1) - ) * _positive_int(config.get("gradient_accumulation_steps"), 1) - previous_rows = (step * per_step) / epoch - # Twice the bound, so an eval carve or masked-out rows in the earlier leg - # cannot read as a different dataset. - return previous_rows > max_train_rows * 2 + return None, fallback_seed + return _positive_int(recorded, 0) or None, _seed_int(marker.get("seed"), fallback_seed) def bound_dataset_rows( diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 7f353a357..fb7c9b201 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -89,15 +89,7 @@ should_use_mlx_training_backend, ) -# Re-exported: the worker and the tests import these from here, and the MLX -# loader imports them from the light module directly. -from .dataset_bounds import ( # noqa: F401 - MAX_STEPS_ROW_SLACK, - MIN_MAX_STEPS_ROWS, - bound_dataset_rows, - max_steps_dataset_rows, - max_train_rows_for_config, -) +from .dataset_bounds import bound_dataset_rows logger = get_logger(__name__) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index e93dbd467..b57f717ac 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -59,8 +59,9 @@ # so it cannot reach these through core.training.trainer. from core.training.dataset_bounds import ( bound_dataset_rows, - checkpoint_predates_row_bound, max_train_rows_for_config, + record_row_bound, + row_bound_for_resume, ) from utils.training_runs import build_default_output_dir_name from utils.wheel_utils import ( @@ -2792,6 +2793,13 @@ def _send(event_type, **kwargs): # Recomputed from the config rather than carried over from the parent so an # MLX run can never train against a bound derived from stale values. mlx_max_train_rows = max_train_rows_for_config(config, is_vlm = is_vlm) + # MLXTrainer resumes by jumping a batch cursor into a schedule rebuilt from + # whatever dataset it is handed, so a bound applied to a checkpoint that was + # written without one continues on unrelated rows. Same marker, same rule as + # the CUDA path. + mlx_max_train_rows, mlx_max_train_rows_seed = row_bound_for_resume( + resume_from_checkpoint, mlx_max_train_rows, random_seed + ) def _slice(ds): if slice_start is not None or slice_end is not None: @@ -2804,7 +2812,7 @@ def _slice(ds): return bound_dataset_rows( ds, mlx_max_train_rows, - config.get("random_seed", 3407), + mlx_max_train_rows_seed, on_bound = lambda kept, total: _send( "status", status_message = f"Using {kept} of {total} rows (max_steps run)", @@ -3016,6 +3024,8 @@ def _fmt_progress(status_message = "", **_kw): ) ensure_dir(Path(output_dir)) _emit_output_dir(event_queue, output_dir) + # Pin the subset before any checkpoint lands here; a resume reads it back. + record_row_bound(output_dir, mlx_max_train_rows, mlx_max_train_rows_seed) # ── 6. Create trainer ── raw_eval_steps = config.get("eval_steps", 0) @@ -4112,19 +4122,21 @@ def _apply_stop(save: bool) -> None: # explicit train-split range opt out inside load_and_format_dataset, where # they live. max_train_rows = max_train_rows_for_config(config) - # A resume keeps the same rows, because the bound is a function of the - # seed, max_steps, batch size and accumulation. A checkpoint from before - # the bound existed is the exception: it trained on the whole dataset, and - # the trainer fast-forwards by batch count over the current dataloader, so - # bounding it now would continue into unrelated rows. - if checkpoint_predates_row_bound( - config.get("resume_from_checkpoint"), max_train_rows, config - ): + max_train_rows_seed = config.get("random_seed", 3407) + # A resume trains on the rows its first start chose, read back from the + # marker written beside the checkpoints. A checkpoint with no marker + # predates the bound: it trained on the whole dataset, and the trainer + # fast-forwards by batch count over the current dataloader, so bounding it + # now would continue into unrelated rows. + resumed_rows, max_train_rows_seed = row_bound_for_resume( + config.get("resume_from_checkpoint"), max_train_rows, max_train_rows_seed + ) + if resumed_rows != max_train_rows: logger.info( - "Resuming a checkpoint that trained on the full dataset: skipping the " - "max_steps row bound so the run continues over the same rows\n" + "Resuming with the row bound recorded at the original start " + f"({resumed_rows} rows) instead of {max_train_rows}\n" ) - max_train_rows = None + max_train_rows = resumed_rows def _load_training_dataset(): result = trainer.load_and_format_dataset( @@ -4150,7 +4162,7 @@ def _load_training_dataset(): or config.get("require_exact_dataset_resource") ), max_train_rows = max_train_rows, - max_train_rows_seed = config.get("random_seed", 3407), + max_train_rows_seed = max_train_rows_seed, ) if isinstance(result, tuple): loaded_dataset, loaded_eval_dataset = result @@ -4507,6 +4519,10 @@ def _monitor_tqdm(): output_dir = str(resolve_output_dir(output_dir)) ensure_dir(Path(output_dir)) _emit_output_dir(event_queue, output_dir) + # Pin the subset this run trains on before any checkpoint lands in here, + # so a later resume reads it back rather than deriving it from a config + # the user may have edited in between. + record_row_bound(output_dir, max_train_rows, max_train_rows_seed) tensorboard_dir = config.get("tensorboard_dir") if config.get("enable_tensorboard", False): diff --git a/studio/backend/tests/test_training_preflight.py b/studio/backend/tests/test_training_preflight.py index 19876f868..ab2c83fcd 100644 --- a/studio/backend/tests/test_training_preflight.py +++ b/studio/backend/tests/test_training_preflight.py @@ -504,7 +504,7 @@ def fail_remote(*args, **kwargs): def test_max_steps_dataset_rows_bounds_the_run(): - from core.training.trainer import ( + from core.training.dataset_bounds import ( MAX_STEPS_ROW_SLACK, MIN_MAX_STEPS_ROWS, max_steps_dataset_rows, @@ -589,7 +589,7 @@ def test_max_steps_bound_is_off_without_it(monkeypatch): def test_max_steps_dataset_rows_survives_unusable_numbers(): - from core.training.trainer import MIN_MAX_STEPS_ROWS, max_steps_dataset_rows + from core.training.dataset_bounds import MIN_MAX_STEPS_ROWS, max_steps_dataset_rows # The worker is also driven from the DB and by direct callers, so a None or a # string reaches this. A row bound is an optimization; it must never raise. @@ -704,48 +704,59 @@ def test_bound_dataset_rows_leaves_a_dataset_dict_alone(): assert bound_dataset_rows(splits, 1024, 3407) is splits -def test_checkpoint_predates_row_bound(tmp_path): - import json +def test_row_bound_marker_round_trips_through_a_resume(tmp_path): + from core.training.dataset_bounds import record_row_bound, row_bound_for_resume + + run_dir = tmp_path / "run" + run_dir.mkdir() + checkpoint = run_dir / "checkpoint-30" + checkpoint.mkdir() + + # Not resuming: the freshly computed pair, which is what gets recorded. + assert row_bound_for_resume(None, 4096, 3407) == (4096, 3407) - from core.training.dataset_bounds import checkpoint_predates_row_bound + record_row_bound(str(run_dir), 4096, 3407) + # Resuming reads back the bound the run started with, so the rows and their + # order do not move when max_steps or the batch size are edited in between. + assert row_bound_for_resume(str(checkpoint), 40960, 99) == (4096, 3407) + # The run directory itself is accepted as well as a checkpoint inside it. + assert row_bound_for_resume(str(run_dir), 40960, 99) == (4096, 3407) - config = {"batch_size": 2, "gradient_accumulation_steps": 4} + # A run that was never bounded stays unbounded on resume. + unbounded = tmp_path / "unbounded" + unbounded.mkdir() + record_row_bound(str(unbounded), None, 3407) + assert row_bound_for_resume(str(unbounded / "checkpoint-5"), 1024, 3407) == (None, 3407) - def _checkpoint(name, *, step, epoch): - path = tmp_path / name - path.mkdir() - (path / "trainer_state.json").write_text(json.dumps({"global_step": step, "epoch": epoch})) - return str(path) - # 15 steps x 8 rows against 192,523 rows: a run that saw the whole corpus. - full = _checkpoint("full", step = 15, epoch = 120 / 192_523) - assert checkpoint_predates_row_bound(full, 1024, config) is True +def test_row_bound_is_dropped_for_a_checkpoint_that_predates_it(tmp_path): + from core.training.dataset_bounds import record_row_bound, row_bound_for_resume - # The same 15 steps against a 1,024-row bound: already bounded, so resuming - # keeps the bound and continues over the same rows. - bounded = _checkpoint("bounded", step = 15, epoch = 120 / 1024) - assert checkpoint_predates_row_bound(bounded, 1024, config) is False + # A checkpoint written before the marker existed trained on the whole corpus + # in its natural order. Both trainers resume by batch index, so a shuffled + # subset would continue on unrelated rows: no bound, whatever its size. + legacy = tmp_path / "legacy" + (legacy / "checkpoint-30").mkdir(parents = True) + assert row_bound_for_resume(str(legacy / "checkpoint-30"), 1024, 3407) == (None, 3407) - # The checkpoint's own batch size wins over a resume that changed it. - recorded = tmp_path / "recorded" - recorded.mkdir() - (recorded / "trainer_state.json").write_text( - json.dumps({"global_step": 15, "epoch": 120 / 192_523, "train_batch_size": 2}) + # Including the range the arithmetic estimate could not tell apart: a legacy + # dataset only slightly larger than the bound still gets shrunk by it. + (legacy / "checkpoint-30" / "trainer_state.json").write_text( + json.dumps({"global_step": 15, "epoch": 120 / 1500, "train_batch_size": 2}) ) - assert checkpoint_predates_row_bound(str(recorded), 1024, {**config, "batch_size": 8}) is True - - # Nothing to decide from, and nothing to break: leave the bound alone. - assert checkpoint_predates_row_bound(None, 1024, config) is False - assert checkpoint_predates_row_bound(full, None, config) is False - assert checkpoint_predates_row_bound(str(tmp_path / "missing"), 1024, config) is False - empty = tmp_path / "empty" - empty.mkdir() - (empty / "trainer_state.json").write_text("{}") - assert checkpoint_predates_row_bound(str(empty), 1024, config) is False - broken = tmp_path / "broken" - broken.mkdir() - (broken / "trainer_state.json").write_text("not json") - assert checkpoint_predates_row_bound(str(broken), 1024, config) is False + assert row_bound_for_resume(str(legacy / "checkpoint-30"), 1024, 3407) == (None, 3407) + + # An unreadable or truncated marker reads as legacy, never as a bound. + for name, body in (("empty", "{}"), ("broken", "not json"), ("null", "null")): + run_dir = tmp_path / name + run_dir.mkdir() + (run_dir / "unsloth_row_bound.json").write_text(body) + assert row_bound_for_resume(str(run_dir), 1024, 3407) == (None, 3407) + + # A marker that cannot be written leaves the resume unbounded rather than + # failing the run that was trying to record it. + record_row_bound(str(tmp_path / "does" / "not" / "exist"), 1024, 3407) + record_row_bound(None, 1024, 3407) def test_bound_dataset_rows_is_deterministic_and_seed_sensitive(): @@ -815,6 +826,11 @@ def test_both_loaders_apply_the_row_bound(): # The MLX worker loads its own dataset, so it has to bound its own rows. assert "bound_dataset_rows" in calls["_slice"] assert "max_train_rows_for_config" in calls["_run_mlx_training"] + # Both loaders resume on the recorded bound and both record their own, or a + # resume silently trains on rows the checkpoint never saw. + for loader in ("run_training_process", "_run_mlx_training"): + assert "row_bound_for_resume" in calls[loader] + assert "record_row_bound" in calls[loader] def test_mlx_adapter_keeps_one_source_of_truth_for_the_bound(): From dcf8187a2762edfeec3b0d57e775f533738b95d2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 15 Aug 2026 14:53:51 +0000 Subject: [PATCH 06/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/training/dataset_bounds.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/studio/backend/core/training/dataset_bounds.py b/studio/backend/core/training/dataset_bounds.py index 8e214548f..beda7b054 100644 --- a/studio/backend/core/training/dataset_bounds.py +++ b/studio/backend/core/training/dataset_bounds.py @@ -120,7 +120,11 @@ def run_dir_for_checkpoint(checkpoint_path: Any) -> Optional[str]: return path -def record_row_bound(output_dir: Any, max_train_rows: Optional[int], seed: Any = 3407) -> None: +def record_row_bound( + output_dir: Any, + max_train_rows: Optional[int], + seed: Any = 3407, +) -> None: """Record the bound a run started with, beside its checkpoints. The subset a run trains on is training state: it has to be fixed at the first @@ -138,9 +142,7 @@ def record_row_bound(output_dir: Any, max_train_rows: Optional[int], seed: Any = if not run_dir: return try: - with open( - os.path.join(run_dir, ROW_BOUND_MARKER_FILE), "w", encoding = "utf-8" - ) as handle: + with open(os.path.join(run_dir, ROW_BOUND_MARKER_FILE), "w", encoding = "utf-8") as handle: json.dump( { "max_train_rows": _positive_int(max_train_rows, 0) or None, @@ -153,7 +155,9 @@ def record_row_bound(output_dir: Any, max_train_rows: Optional[int], seed: Any = def row_bound_for_resume( - checkpoint_path: Any, max_train_rows: Optional[int], seed: Any = 3407 + checkpoint_path: Any, + max_train_rows: Optional[int], + seed: Any = 3407, ) -> tuple[Optional[int], int]: """The (rows, seed) a resume must use, or the freshly computed pair. @@ -179,9 +183,7 @@ def row_bound_for_resume( if not run_dir: return max_train_rows, fallback_seed try: - with open( - os.path.join(run_dir, ROW_BOUND_MARKER_FILE), encoding = "utf-8" - ) as handle: + with open(os.path.join(run_dir, ROW_BOUND_MARKER_FILE), encoding = "utf-8") as handle: marker = json.load(handle) recorded = marker["max_train_rows"] except (OSError, UnicodeDecodeError, ValueError, TypeError, KeyError): From 06d1a5e4cc62337263f586d616882b6a1c35b31b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 15 Aug 2026 15:06:58 +0000 Subject: [PATCH 07/14] harden the row-bound marker: atomic write, exact checkpoint match, raw-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- 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. --- .../backend/core/training/dataset_bounds.py | 63 +++++++++++++++---- .../backend/tests/test_training_preflight.py | 54 ++++++++++++++++ 2 files changed, 105 insertions(+), 12 deletions(-) diff --git a/studio/backend/core/training/dataset_bounds.py b/studio/backend/core/training/dataset_bounds.py index beda7b054..d2147ff64 100644 --- a/studio/backend/core/training/dataset_bounds.py +++ b/studio/backend/core/training/dataset_bounds.py @@ -14,6 +14,8 @@ import json import os +import re +import tempfile from typing import Any, Optional # Slack on the row bound. Rows are consumed by things that never produce a step: @@ -28,6 +30,8 @@ # Written into a run's output directory at its first start; read back on resume. # Its absence is the signal that a checkpoint predates the bound. ROW_BOUND_MARKER_FILE = "unsloth_row_bound.json" +# transformers writes checkpoint- and nothing else under that prefix. +_CHECKPOINT_DIR_RE = re.compile(r"^checkpoint-\d+$") def _int_or(value: Any, default: int) -> int: @@ -78,11 +82,23 @@ def effective_packing(config: dict, is_vlm: bool = False) -> bool: Packing opts the bound out because one packed sample spans an unknown number of source rows. The stored value alone overshoots: the frontend hides the packing control for image VLMs without resetting it, API clients can submit - the combination directly, and the image, audio-codec and audio-VLM branches - all train without packing whatever the config says. + the combination directly, and the vision, audio-VLM and audio-codec branches + all train without packing whatever the config says -- csm and snac on a plain + HF Trainer, whisper on Seq2SeqTrainer, none of which take a packing argument, + and bicodec/dac forced off on the SFTTrainer path. + + Raw-text and CPT are the exception to that: the vision/audio-VLM branch is + gated on `not raw_text_mode`, so those runs take the text path and it honours + the requested value however the dataset is flagged. """ if not config.get("packing", False): return False + raw_text_mode = ( + config.get("training_type") == "Continued Pretraining" + or config.get("format_type") == "raw" + ) + if raw_text_mode: + return True if is_vlm or config.get("is_dataset_image", False) or config.get("is_dataset_audio", False): return False return True @@ -106,7 +122,11 @@ def max_train_rows_for_config(config: dict, is_vlm: bool = False) -> Optional[in def run_dir_for_checkpoint(checkpoint_path: Any) -> Optional[str]: """The run directory a checkpoint lives in, or None when there is none. - Checkpoints are written as ``/checkpoint-``; a caller that + Trainer writes ``/checkpoint-`` (transformers' + PREFIX_CHECKPOINT_DIR, always followed by the step number), so only that + exact shape is a checkpoint. Matching the bare prefix would take the parent + of a RUN directory that happens to start with it, and the marker would then + be written one level above where a later resume looks for it. A caller that names the run directory itself gets it back unchanged. """ if not checkpoint_path: @@ -115,7 +135,7 @@ def run_dir_for_checkpoint(checkpoint_path: Any) -> Optional[str]: if not path: return None head, tail = os.path.split(path) - if tail.startswith("checkpoint-") and head: + if head and _CHECKPOINT_DIR_RE.match(tail): return head return path @@ -137,21 +157,40 @@ def record_row_bound( Best effort by design. A run must never fail over a marker; a marker that could not be written costs the optimization on the next resume and nothing else. + + Written through a temporary file and moved into place, because a resume + rewrites a marker that is already valid: truncating in place and then failing + -- a full disk is the ordinary way -- would leave an empty file, which reads + as "no marker" and resumes the run over the whole dataset. os.replace is + atomic on POSIX and on Windows. """ run_dir = run_dir_for_checkpoint(output_dir) if not run_dir: return + marker = os.path.join(run_dir, ROW_BOUND_MARKER_FILE) + tmp_path = None try: - with open(os.path.join(run_dir, ROW_BOUND_MARKER_FILE), "w", encoding = "utf-8") as handle: - json.dump( - { - "max_train_rows": _positive_int(max_train_rows, 0) or None, - "seed": _seed_int(seed, 3407), - }, - handle, - ) + payload = json.dumps( + { + "max_train_rows": _positive_int(max_train_rows, 0) or None, + "seed": _seed_int(seed, 3407), + } + ) + handle, tmp_path = tempfile.mkstemp(dir = run_dir, prefix = ".row_bound_", suffix = ".tmp") + with os.fdopen(handle, "w", encoding = "utf-8") as tmp_file: + tmp_file.write(payload) + tmp_file.flush() + os.fsync(tmp_file.fileno()) + os.replace(tmp_path, marker) + tmp_path = None except (OSError, UnicodeError, TypeError, ValueError): return + finally: + if tmp_path is not None: + try: + os.unlink(tmp_path) + except OSError: + pass def row_bound_for_resume( diff --git a/studio/backend/tests/test_training_preflight.py b/studio/backend/tests/test_training_preflight.py index ab2c83fcd..e5697d99c 100644 --- a/studio/backend/tests/test_training_preflight.py +++ b/studio/backend/tests/test_training_preflight.py @@ -623,6 +623,17 @@ def test_effective_packing_decides_the_opt_out(): # An epoch-bounded run is unbounded whatever packing says. assert max_train_rows_for_config({"max_steps": 0, "packing": False}) is None + # Raw-text and CPT take the text path however the dataset is flagged -- the + # vision branch is gated on `not raw_text_mode` -- and that path honours the + # requested value, so those runs pack for real and keep the opt-out. + for raw in ({"training_type": "Continued Pretraining"}, {"format_type": "raw"}): + assert effective_packing({**text, **raw, "packing": True, "is_dataset_image": True}) is True + assert effective_packing({**text, **raw, "packing": True, "is_dataset_audio": True}) is True + assert max_train_rows_for_config({**text, **raw, "packing": True}) is None + # Without packing they are bounded like anything else. + assert effective_packing({**text, **raw, "is_dataset_image": True}) is False + assert max_train_rows_for_config({**text, **raw, "is_dataset_image": True}) == 1024 + def test_bound_dataset_rows_edges(): from core.training.dataset_bounds import bound_dataset_rows @@ -729,6 +740,49 @@ def test_row_bound_marker_round_trips_through_a_resume(tmp_path): assert row_bound_for_resume(str(unbounded / "checkpoint-5"), 1024, 3407) == (None, 3407) +def test_row_bound_marker_survives_a_run_directory_named_like_a_checkpoint(tmp_path): + from core.training.dataset_bounds import record_row_bound, row_bound_for_resume + + # Trainer writes checkpoint-. A run directory whose own name + # merely starts with the prefix is not one, and taking its parent would file + # the marker one level above where the resume looks for it. + run_dir = tmp_path / "checkpoint-model__project-x" + (run_dir / "checkpoint-30").mkdir(parents = True) + record_row_bound(str(run_dir), 4096, 3407) + + assert (run_dir / "unsloth_row_bound.json").exists() + assert not (tmp_path / "unsloth_row_bound.json").exists() + assert row_bound_for_resume(str(run_dir / "checkpoint-30"), 40960, 99) == (4096, 3407) + + +def test_row_bound_marker_is_replaced_atomically(tmp_path): + import os + + from core.training.dataset_bounds import record_row_bound, row_bound_for_resume + + # A resume rewrites a marker that is already valid. Truncating in place and + # then failing -- a full disk is the ordinary way -- would leave an empty + # file, which reads as "no marker" and resumes over the whole dataset. + run_dir = tmp_path / "run" + run_dir.mkdir() + record_row_bound(str(run_dir), 4096, 3407) + + real_replace = os.replace + + def _fail_replace(src, dst): + raise OSError(28, "No space left on device") + + os.replace = _fail_replace + try: + record_row_bound(str(run_dir), 8192, 99) + finally: + os.replace = real_replace + + assert row_bound_for_resume(str(run_dir), 40960, 99) == (4096, 3407) + # And the temporary file it wrote instead is cleaned up. + assert [p.name for p in run_dir.iterdir()] == ["unsloth_row_bound.json"] + + def test_row_bound_is_dropped_for_a_checkpoint_that_predates_it(tmp_path): from core.training.dataset_bounds import record_row_bound, row_bound_for_resume From a75031e00dcd33a857e6f48b5e37f0d286f95d48 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 15 Aug 2026 15:10:31 +0000 Subject: [PATCH 08/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/training/dataset_bounds.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/studio/backend/core/training/dataset_bounds.py b/studio/backend/core/training/dataset_bounds.py index d2147ff64..617243020 100644 --- a/studio/backend/core/training/dataset_bounds.py +++ b/studio/backend/core/training/dataset_bounds.py @@ -94,8 +94,7 @@ def effective_packing(config: dict, is_vlm: bool = False) -> bool: if not config.get("packing", False): return False raw_text_mode = ( - config.get("training_type") == "Continued Pretraining" - or config.get("format_type") == "raw" + config.get("training_type") == "Continued Pretraining" or config.get("format_type") == "raw" ) if raw_text_mode: return True From 5b4ec486e48295b4571feeb68afd4265482ac2c1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 15 Aug 2026 15:22:21 +0000 Subject: [PATCH 09/14] stop guessing the training branch from the dataset flags 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. --- .../backend/core/training/dataset_bounds.py | 43 ++++++++++--------- studio/backend/core/training/worker.py | 20 ++++++++- .../backend/tests/test_training_preflight.py | 39 +++++++++++------ 3 files changed, 65 insertions(+), 37 deletions(-) diff --git a/studio/backend/core/training/dataset_bounds.py b/studio/backend/core/training/dataset_bounds.py index 617243020..42fdd5208 100644 --- a/studio/backend/core/training/dataset_bounds.py +++ b/studio/backend/core/training/dataset_bounds.py @@ -80,27 +80,25 @@ def effective_packing(config: dict, is_vlm: bool = False) -> bool: """Whether the trainer will actually pack, not merely what was requested. Packing opts the bound out because one packed sample spans an unknown number - of source rows. The stored value alone overshoots: the frontend hides the - packing control for image VLMs without resetting it, API clients can submit - the combination directly, and the vision, audio-VLM and audio-codec branches - all train without packing whatever the config says -- csm and snac on a plain - HF Trainer, whisper on Seq2SeqTrainer, none of which take a packing argument, - and bicodec/dac forced off on the SFTTrainer path. - - Raw-text and CPT are the exception to that: the vision/audio-VLM branch is - gated on `not raw_text_mode`, so those runs take the text path and it honours - the requested value however the dataset is flagged. + of source rows, and the requested value is the answer unless the caller has + established that the run cannot pack. Only `is_vlm` does that: it says the + caller probed the model and the dataset and landed on the vision branch, + which sets no packing at all. + + The dataset flags do NOT establish it, though they look like they should. + `is_dataset_image` and `is_dataset_audio` are client-supplied and true on a + column-NAME match, so a text model with a column called "audio" carries the + flag and still trains on the text path, which honours packing. Raw-text and + CPT take that path too, whatever the flags say, because the vision branch is + gated on `not raw_text_mode`. Guessing the branch from a flag was wrong in + three separate ways; where the branch is unknown, assume it packs. """ if not config.get("packing", False): return False raw_text_mode = ( config.get("training_type") == "Continued Pretraining" or config.get("format_type") == "raw" ) - if raw_text_mode: - return True - if is_vlm or config.get("is_dataset_image", False) or config.get("is_dataset_audio", False): - return False - return True + return bool(raw_text_mode or not is_vlm) def max_train_rows_for_config(config: dict, is_vlm: bool = False) -> Optional[int]: @@ -143,7 +141,7 @@ def record_row_bound( output_dir: Any, max_train_rows: Optional[int], seed: Any = 3407, -) -> None: +) -> bool: """Record the bound a run started with, beside its checkpoints. The subset a run trains on is training state: it has to be fixed at the first @@ -153,9 +151,11 @@ def record_row_bound( runs, and a checkpoint written before this feature existed leaves no arithmetic that distinguishes it reliably. - Best effort by design. A run must never fail over a marker; a marker that - could not be written costs the optimization on the next resume and nothing - else. + Best effort by design, and it answers whether it succeeded so the caller can + say so. A run must never fail over a marker, and by the time this is called + the dataset has already been bounded, so there is nothing to fall back to + either; what an unwritable marker costs is a later resume reading the run as + unbounded. Written through a temporary file and moved into place, because a resume rewrites a marker that is already valid: truncating in place and then failing @@ -165,7 +165,7 @@ def record_row_bound( """ run_dir = run_dir_for_checkpoint(output_dir) if not run_dir: - return + return False marker = os.path.join(run_dir, ROW_BOUND_MARKER_FILE) tmp_path = None try: @@ -183,13 +183,14 @@ def record_row_bound( os.replace(tmp_path, marker) tmp_path = None except (OSError, UnicodeError, TypeError, ValueError): - return + return False finally: if tmp_path is not None: try: os.unlink(tmp_path) except OSError: pass + return True def row_bound_for_resume( diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index b57f717ac..26f7d2d31 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -3025,7 +3025,16 @@ def _fmt_progress(status_message = "", **_kw): ensure_dir(Path(output_dir)) _emit_output_dir(event_queue, output_dir) # Pin the subset before any checkpoint lands here; a resume reads it back. - record_row_bound(output_dir, mlx_max_train_rows, mlx_max_train_rows_seed) + if not record_row_bound(output_dir, mlx_max_train_rows, mlx_max_train_rows_seed) and ( + mlx_max_train_rows + ): + _send( + "warning", + message = ( + f"Could not record the max_steps row bound in {output_dir}: " + "resuming this run later will read it as unbounded" + ), + ) # ── 6. Create trainer ── raw_eval_steps = config.get("eval_steps", 0) @@ -4522,7 +4531,14 @@ def _monitor_tqdm(): # Pin the subset this run trains on before any checkpoint lands in here, # so a later resume reads it back rather than deriving it from a config # the user may have edited in between. - record_row_bound(output_dir, max_train_rows, max_train_rows_seed) + if not record_row_bound(output_dir, max_train_rows, max_train_rows_seed) and max_train_rows: + # Not fatal, and nothing to fall back to at this point: the dataset is + # already bounded. Say it, so a later resume reading this run as + # unbounded is explainable. + logger.warning( + f"Could not record the max_steps row bound in {output_dir}: " + "resuming this run later will read it as unbounded\n" + ) tensorboard_dir = config.get("tensorboard_dir") if config.get("enable_tensorboard", False): diff --git a/studio/backend/tests/test_training_preflight.py b/studio/backend/tests/test_training_preflight.py index e5697d99c..bc91c359f 100644 --- a/studio/backend/tests/test_training_preflight.py +++ b/studio/backend/tests/test_training_preflight.py @@ -611,28 +611,30 @@ def test_effective_packing_decides_the_opt_out(): assert effective_packing({**text, "packing": True}) is True assert max_train_rows_for_config({**text, "packing": True}) is None - # ...but the image, audio and VLM branches train without packing whatever the - # config says, and the frontend hides the control without resetting it. A - # stale flag must not cost those runs the bound. - assert effective_packing({**text, "packing": True, "is_dataset_image": True}) is False - assert effective_packing({**text, "packing": True, "is_dataset_audio": True}) is False + # A caller that probed the model and landed on the vision branch knows the run + # cannot pack, so a stale flag does not cost it the bound. assert effective_packing({**text, "packing": True}, is_vlm = True) is False assert max_train_rows_for_config({**text, "packing": True}, is_vlm = True) == 1024 - assert max_train_rows_for_config({**text, "packing": True, "is_dataset_image": True}) == 1024 + + # The dataset flags alone establish nothing: they are client-supplied and true + # on a column-NAME match, so a text model with a column called "audio" or + # "image" still trains on the text path, which honours packing. + assert effective_packing({**text, "packing": True, "is_dataset_image": True}) is True + assert effective_packing({**text, "packing": True, "is_dataset_audio": True}) is True + assert max_train_rows_for_config({**text, "packing": True, "is_dataset_audio": True}) is None # An epoch-bounded run is unbounded whatever packing says. assert max_train_rows_for_config({"max_steps": 0, "packing": False}) is None - # Raw-text and CPT take the text path however the dataset is flagged -- the - # vision branch is gated on `not raw_text_mode` -- and that path honours the - # requested value, so those runs pack for real and keep the opt-out. + # Raw-text and CPT take the text path even on a vision model -- that branch + # is gated on `not raw_text_mode` -- and the text path honours the requested + # value, so those runs pack for real and keep the opt-out. for raw in ({"training_type": "Continued Pretraining"}, {"format_type": "raw"}): - assert effective_packing({**text, **raw, "packing": True, "is_dataset_image": True}) is True - assert effective_packing({**text, **raw, "packing": True, "is_dataset_audio": True}) is True - assert max_train_rows_for_config({**text, **raw, "packing": True}) is None + assert effective_packing({**text, **raw, "packing": True}, is_vlm = True) is True + assert max_train_rows_for_config({**text, **raw, "packing": True}, is_vlm = True) is None # Without packing they are bounded like anything else. - assert effective_packing({**text, **raw, "is_dataset_image": True}) is False - assert max_train_rows_for_config({**text, **raw, "is_dataset_image": True}) == 1024 + assert effective_packing({**text, **raw}, is_vlm = True) is False + assert max_train_rows_for_config({**text, **raw}, is_vlm = True) == 1024 def test_bound_dataset_rows_edges(): @@ -783,6 +785,15 @@ def _fail_replace(src, dst): assert [p.name for p in run_dir.iterdir()] == ["unsloth_row_bound.json"] +def test_record_row_bound_reports_whether_it_wrote(): + from core.training.dataset_bounds import record_row_bound + + # The caller logs a failure rather than failing the run: by the time this is + # called the dataset is already bounded, so there is nothing to fall back to. + assert record_row_bound(None, 1024, 3407) is False + assert record_row_bound("/definitely/not/a/directory/here", 1024, 3407) is False + + def test_row_bound_is_dropped_for_a_checkpoint_that_predates_it(tmp_path): from core.training.dataset_bounds import record_row_bound, row_bound_for_resume From 9dfff03d97d4bf03c55c9bddb028e37c3e38eb99 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 15 Aug 2026 15:41:45 +0000 Subject: [PATCH 10/14] read the training branch the CUDA worker already detected 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. --- .../backend/core/training/dataset_bounds.py | 31 ++++++------ studio/backend/core/training/worker.py | 49 +++++++++++-------- .../backend/tests/test_training_preflight.py | 43 +++++++++++++--- 3 files changed, 83 insertions(+), 40 deletions(-) diff --git a/studio/backend/core/training/dataset_bounds.py b/studio/backend/core/training/dataset_bounds.py index 42fdd5208..a2d8db201 100644 --- a/studio/backend/core/training/dataset_bounds.py +++ b/studio/backend/core/training/dataset_bounds.py @@ -76,38 +76,39 @@ def max_steps_dataset_rows( return max(MIN_MAX_STEPS_ROWS, steps * per_step * MAX_STEPS_ROW_SLACK) -def effective_packing(config: dict, is_vlm: bool = False) -> bool: +def effective_packing(config: dict, branch_never_packs: bool = False) -> bool: """Whether the trainer will actually pack, not merely what was requested. Packing opts the bound out because one packed sample spans an unknown number of source rows, and the requested value is the answer unless the caller has - established that the run cannot pack. Only `is_vlm` does that: it says the - caller probed the model and the dataset and landed on the vision branch, - which sets no packing at all. + established that the branch this run takes sets no packing at all: the + vision and audio-VLM branches, and every audio codec, which train on a + Trainer that has no packing argument to give. - The dataset flags do NOT establish it, though they look like they should. + The dataset flags do NOT establish that, though they look like they should. `is_dataset_image` and `is_dataset_audio` are client-supplied and true on a column-NAME match, so a text model with a column called "audio" carries the - flag and still trains on the text path, which honours packing. Raw-text and - CPT take that path too, whatever the flags say, because the vision branch is - gated on `not raw_text_mode`. Guessing the branch from a flag was wrong in - three separate ways; where the branch is unknown, assume it packs. + flag and still trains on the text path, which honours packing. Pass the + branch the model probe actually detected instead. + + Raw-text and CPT are the exception on top: they reach the text path from any + branch, because the vision one is gated on `not raw_text_mode`. """ if not config.get("packing", False): return False raw_text_mode = ( config.get("training_type") == "Continued Pretraining" or config.get("format_type") == "raw" ) - return bool(raw_text_mode or not is_vlm) + return bool(raw_text_mode or not branch_never_packs) -def max_train_rows_for_config(config: dict, is_vlm: bool = False) -> Optional[int]: +def max_train_rows_for_config(config: dict, branch_never_packs: bool = False) -> Optional[int]: """The bound for a worker config, or None when the run is not bounded. Streaming and an explicit train-split range opt out further down, in the loaders, where those values live. """ - if effective_packing(config, is_vlm = is_vlm): + if effective_packing(config, branch_never_packs = branch_never_packs): return None return max_steps_dataset_rows( config.get("max_steps", 0) or 0, @@ -132,8 +133,10 @@ def run_dir_for_checkpoint(checkpoint_path: Any) -> Optional[str]: if not path: return None head, tail = os.path.split(path) - if head and _CHECKPOINT_DIR_RE.match(tail): - return head + if _CHECKPOINT_DIR_RE.match(tail): + # A bare "checkpoint-30" splits to an empty head, and its run directory is + # the working directory rather than itself. + return head or os.curdir return path diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 26f7d2d31..3cb78444d 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -2792,7 +2792,7 @@ def _send(event_type, **kwargs): # point -- formatting, chat templating, tokenization -- maps over every row. # Recomputed from the config rather than carried over from the parent so an # MLX run can never train against a bound derived from stale values. - mlx_max_train_rows = max_train_rows_for_config(config, is_vlm = is_vlm) + mlx_max_train_rows = max_train_rows_for_config(config, branch_never_packs = is_vlm) # MLXTrainer resumes by jumping a batch cursor into a schedule rebuilt from # whatever dataset it is handed, so a bound applied to a checkpoint that was # written without one continues on unrelated rows. Same marker, same rule as @@ -4126,26 +4126,10 @@ def _apply_stop(save: bool) -> None: training_type = config.get("training_type", "LoRA/QLoRA") is_cpt_for_dataset = training_type == "Continued Pretraining" - # Effective packing opts out: one packed sample spans an unknown number of - # rows, so steps cannot be converted to a row count. Streaming and an - # explicit train-split range opt out inside load_and_format_dataset, where - # they live. - max_train_rows = max_train_rows_for_config(config) + # Filled in below, after the model probe: the closure reads them when it + # runs, which is after both. + max_train_rows = None max_train_rows_seed = config.get("random_seed", 3407) - # A resume trains on the rows its first start chose, read back from the - # marker written beside the checkpoints. A checkpoint with no marker - # predates the bound: it trained on the whole dataset, and the trainer - # fast-forwards by batch count over the current dataloader, so bounding it - # now would continue into unrelated rows. - resumed_rows, max_train_rows_seed = row_bound_for_resume( - config.get("resume_from_checkpoint"), max_train_rows, max_train_rows_seed - ) - if resumed_rows != max_train_rows: - logger.info( - "Resuming with the row bound recorded at the original start " - f"({resumed_rows} rows) instead of {max_train_rows}\n" - ) - max_train_rows = resumed_rows def _load_training_dataset(): result = trainer.load_and_format_dataset( @@ -4230,6 +4214,31 @@ def _load_training_dataset(): event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()}) return + # Now that 4a has probed the model, the branch this run takes is known, so + # the packing opt-out can read it instead of guessing from the client's + # dataset flags. Streaming and an explicit train-split range opt out inside + # load_and_format_dataset, where they live. + branch_never_packs = bool( + getattr(trainer, "is_vlm", False) + or getattr(trainer, "is_audio_vlm", False) + or getattr(trainer, "_audio_type", None) + ) + max_train_rows = max_train_rows_for_config(config, branch_never_packs = branch_never_packs) + # A resume trains on the rows its first start chose, read back from the + # marker written beside the checkpoints. A checkpoint with no marker + # predates the bound: it trained on the whole dataset, and the trainer + # fast-forwards by batch count over the current dataloader, so bounding it + # now would continue into unrelated rows. + resumed_rows, max_train_rows_seed = row_bound_for_resume( + config.get("resume_from_checkpoint"), max_train_rows, max_train_rows_seed + ) + if resumed_rows != max_train_rows: + logger.info( + "Resuming with the row bound recorded at the original start " + f"({resumed_rows} rows) instead of {max_train_rows}\n" + ) + max_train_rows = resumed_rows + # ── 4b. Load and format dataset (LLM helper may use VRAM briefly) ── _send_status(event_queue, "Loading and formatting dataset...") dataset, eval_dataset = _load_training_dataset() diff --git a/studio/backend/tests/test_training_preflight.py b/studio/backend/tests/test_training_preflight.py index bc91c359f..ae8fe0106 100644 --- a/studio/backend/tests/test_training_preflight.py +++ b/studio/backend/tests/test_training_preflight.py @@ -613,8 +613,8 @@ def test_effective_packing_decides_the_opt_out(): # A caller that probed the model and landed on the vision branch knows the run # cannot pack, so a stale flag does not cost it the bound. - assert effective_packing({**text, "packing": True}, is_vlm = True) is False - assert max_train_rows_for_config({**text, "packing": True}, is_vlm = True) == 1024 + assert effective_packing({**text, "packing": True}, branch_never_packs = True) is False + assert max_train_rows_for_config({**text, "packing": True}, branch_never_packs = True) == 1024 # The dataset flags alone establish nothing: they are client-supplied and true # on a column-NAME match, so a text model with a column called "audio" or @@ -630,11 +630,11 @@ def test_effective_packing_decides_the_opt_out(): # is gated on `not raw_text_mode` -- and the text path honours the requested # value, so those runs pack for real and keep the opt-out. for raw in ({"training_type": "Continued Pretraining"}, {"format_type": "raw"}): - assert effective_packing({**text, **raw, "packing": True}, is_vlm = True) is True - assert max_train_rows_for_config({**text, **raw, "packing": True}, is_vlm = True) is None + assert effective_packing({**text, **raw, "packing": True}, branch_never_packs = True) is True + assert max_train_rows_for_config({**text, **raw, "packing": True}, branch_never_packs = True) is None # Without packing they are bounded like anything else. - assert effective_packing({**text, **raw}, is_vlm = True) is False - assert max_train_rows_for_config({**text, **raw}, is_vlm = True) == 1024 + assert effective_packing({**text, **raw}, branch_never_packs = True) is False + assert max_train_rows_for_config({**text, **raw}, branch_never_packs = True) == 1024 def test_bound_dataset_rows_edges(): @@ -785,6 +785,28 @@ def _fail_replace(src, dst): assert [p.name for p in run_dir.iterdir()] == ["unsloth_row_bound.json"] +def test_run_dir_for_a_bare_relative_checkpoint(tmp_path, monkeypatch): + from core.training.dataset_bounds import ( + record_row_bound, + row_bound_for_resume, + run_dir_for_checkpoint, + ) + + # "checkpoint-30" splits to an empty head; its run directory is the working + # directory, not itself, or the marker is looked for one level too deep and + # the bounded run reads as legacy. + assert run_dir_for_checkpoint("checkpoint-30") == os.curdir + assert run_dir_for_checkpoint("run/checkpoint-30") == "run" + # A relative run directory is still itself. + assert run_dir_for_checkpoint("checkpoint-model") == "checkpoint-model" + + run_dir = tmp_path / "run" + (run_dir / "checkpoint-30").mkdir(parents = True) + record_row_bound(str(run_dir), 4096, 3407) + monkeypatch.chdir(run_dir) + assert row_bound_for_resume("checkpoint-30", 40960, 99) == (4096, 3407) + + def test_record_row_bound_reports_whether_it_wrote(): from core.training.dataset_bounds import record_row_bound @@ -897,6 +919,15 @@ def test_both_loaders_apply_the_row_bound(): assert "row_bound_for_resume" in calls[loader] assert "record_row_bound" in calls[loader] + # Both pass the branch they detected, rather than letting it default: the + # dataset flags are client-supplied and cannot stand in for it. + assert worker_src.count("branch_never_packs = ") >= 2 + # And the CUDA one computes it only after the model probe has set it, which is + # what makes the value real. + assert worker_src.index("_pre_detect_training_model(\n") < worker_src.index( + "branch_never_packs = bool(" + ) + def test_mlx_adapter_keeps_one_source_of_truth_for_the_bound(): from core.training.training import _build_training_worker_config From d84e44fe82192696e101e899471b8f04d7594cf4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 15 Aug 2026 15:42:56 +0000 Subject: [PATCH 11/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_training_preflight.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/studio/backend/tests/test_training_preflight.py b/studio/backend/tests/test_training_preflight.py index ae8fe0106..494f89865 100644 --- a/studio/backend/tests/test_training_preflight.py +++ b/studio/backend/tests/test_training_preflight.py @@ -631,7 +631,10 @@ def test_effective_packing_decides_the_opt_out(): # value, so those runs pack for real and keep the opt-out. for raw in ({"training_type": "Continued Pretraining"}, {"format_type": "raw"}): assert effective_packing({**text, **raw, "packing": True}, branch_never_packs = True) is True - assert max_train_rows_for_config({**text, **raw, "packing": True}, branch_never_packs = True) is None + assert ( + max_train_rows_for_config({**text, **raw, "packing": True}, branch_never_packs = True) + is None + ) # Without packing they are bounded like anything else. assert effective_packing({**text, **raw}, branch_never_packs = True) is False assert max_train_rows_for_config({**text, **raw}, branch_never_packs = True) == 1024 From 649a0123cf996f77de0c8da8ab93b873d330cfda Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 15 Aug 2026 16:42:33 +0000 Subject: [PATCH 12/14] defer to a bracketed split, and let audio outrank raw mode 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. --- .../backend/core/training/dataset_bounds.py | 13 +++---- studio/backend/core/training/trainer.py | 4 +++ studio/backend/core/training/worker.py | 26 +++++++++++--- .../backend/tests/test_training_preflight.py | 34 +++++++++++++++---- 4 files changed, 57 insertions(+), 20 deletions(-) diff --git a/studio/backend/core/training/dataset_bounds.py b/studio/backend/core/training/dataset_bounds.py index a2d8db201..e514939d7 100644 --- a/studio/backend/core/training/dataset_bounds.py +++ b/studio/backend/core/training/dataset_bounds.py @@ -89,17 +89,14 @@ def effective_packing(config: dict, branch_never_packs: bool = False) -> bool: `is_dataset_image` and `is_dataset_audio` are client-supplied and true on a column-NAME match, so a text model with a column called "audio" carries the flag and still trains on the text path, which honours packing. Pass the - branch the model probe actually detected instead. - - Raw-text and CPT are the exception on top: they reach the text path from any - branch, because the vision one is gated on `not raw_text_mode`. + branch the model probe actually detected instead, which the caller works out: + the two branches differ on raw-text and CPT, since the vision one is gated on + `not raw_text_mode` while audio preprocessing is chosen before the raw-text + bypass and so holds either way. """ if not config.get("packing", False): return False - 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) + return not branch_never_packs def max_train_rows_for_config(config: dict, branch_never_packs: bool = False) -> Optional[int]: diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index fb7c9b201..67c63405b 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -3028,10 +3028,14 @@ def _apply_raw_text_prep(ds: Dataset, split_name: str) -> Dataset: # over every row: that is the cost this avoids. Skipped when the user named # an explicit range, which is already the rows they asked for, and when # streaming, which was bounded lazily above. + # A bracketed split instruction names rows the same way the numeric + # fields do: train[1000:2000] is the user's selection, not a corpus to + # sample from. if ( (not dataset_streaming) and dataset_slice_start is None and dataset_slice_end is None + and "[" not in (train_split or "") ): def _log_bound(kept, total): diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 3cb78444d..068eeef86 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -2792,7 +2792,14 @@ def _send(event_type, **kwargs): # point -- formatting, chat templating, tokenization -- maps over every row. # Recomputed from the config rather than carried over from the parent so an # MLX run can never train against a bound derived from stale values. - mlx_max_train_rows = max_train_rows_for_config(config, branch_never_packs = is_vlm) + # The vision branch is gated on `not raw_text_mode`, so a raw or CPT run + # takes the text path and that path honours the requested packing. + mlx_raw_text_mode = ( + training_type == "Continued Pretraining" or config.get("format_type") == "raw" + ) + mlx_max_train_rows = max_train_rows_for_config( + config, branch_never_packs = is_vlm and not mlx_raw_text_mode + ) # MLXTrainer resumes by jumping a batch cursor into a schedule rebuilt from # whatever dataset it is handed, so a bound applied to a checkpoint that was # written without one continues on unrelated rows. Same marker, same rule as @@ -2801,6 +2808,9 @@ def _send(event_type, **kwargs): resume_from_checkpoint, mlx_max_train_rows, random_seed ) + # A bracketed split instruction names rows the same way the numeric fields do. + mlx_split_names_rows = "[" in (config.get("train_split") or "") + def _slice(ds): if slice_start is not None or slice_end is not None: start = slice_start if slice_start is not None else 0 @@ -2809,6 +2819,8 @@ def _slice(ds): return ds.select([]) # The user named these rows; the bound below defers to that. return ds.select(range(start, min(end + 1, len(ds)))) + if mlx_split_names_rows: + return ds return bound_dataset_rows( ds, mlx_max_train_rows, @@ -4218,10 +4230,14 @@ def _load_training_dataset(): # the packing opt-out can read it instead of guessing from the client's # dataset flags. Streaming and an explicit train-split range opt out inside # load_and_format_dataset, where they live. - branch_never_packs = bool( - getattr(trainer, "is_vlm", False) - or getattr(trainer, "is_audio_vlm", False) - or getattr(trainer, "_audio_type", None) + # Audio codecs are chosen before the raw-text bypass and use plain + # Trainers with no packing argument, so they hold either way; the vision + # and audio-VLM branches are gated on `not raw_text_mode` and give the + # text path, which honours packing, when the run is raw or CPT. + raw_text_mode = is_cpt_for_dataset or config.get("format_type") == "raw" + branch_never_packs = bool(getattr(trainer, "_audio_type", None)) or ( + bool(getattr(trainer, "is_vlm", False) or getattr(trainer, "is_audio_vlm", False)) + and not raw_text_mode ) max_train_rows = max_train_rows_for_config(config, branch_never_packs = branch_never_packs) # A resume trains on the rows its first start chose, read back from the diff --git a/studio/backend/tests/test_training_preflight.py b/studio/backend/tests/test_training_preflight.py index 494f89865..f79be6b47 100644 --- a/studio/backend/tests/test_training_preflight.py +++ b/studio/backend/tests/test_training_preflight.py @@ -574,6 +574,24 @@ def test_max_steps_bound_defers_to_an_explicit_slice(monkeypatch): assert sliced.shuffle_seeds == [] +def test_max_steps_bound_defers_to_a_split_instruction(monkeypatch): + train = _SizedDataset(500_000) + trainer = _cached_only_loader(monkeypatch, train) + + result = trainer.load_and_format_dataset( + "org/dataset", + dataset_local_files_only = True, + dataset_local_path = "/cache/snapshot", + train_split = "train[1000:200000]", + max_train_rows = 1024, + ) + + assert result is not None + # A bracketed split names rows exactly as the numeric slice fields do, so the + # bound must not resample a selection the user already made. + assert result[0]["dataset"] is train + + def test_max_steps_bound_is_off_without_it(monkeypatch): train = _SizedDataset(500_000) trainer = _cached_only_loader(monkeypatch, train) @@ -626,15 +644,15 @@ def test_effective_packing_decides_the_opt_out(): # An epoch-bounded run is unbounded whatever packing says. assert max_train_rows_for_config({"max_steps": 0, "packing": False}) is None - # Raw-text and CPT take the text path even on a vision model -- that branch - # is gated on `not raw_text_mode` -- and the text path honours the requested - # value, so those runs pack for real and keep the opt-out. + # Raw-text and CPT do not enter into it here: the caller decides the branch, + # because the two differ on raw mode. The vision branch is gated on + # `not raw_text_mode` and gives way to the text path, while audio + # preprocessing is chosen before the raw-text bypass and holds either way. for raw in ({"training_type": "Continued Pretraining"}, {"format_type": "raw"}): - assert effective_packing({**text, **raw, "packing": True}, branch_never_packs = True) is True assert ( - max_train_rows_for_config({**text, **raw, "packing": True}, branch_never_packs = True) - is None + effective_packing({**text, **raw, "packing": True}, branch_never_packs = True) is False ) + assert effective_packing({**text, **raw, "packing": True}) is True # Without packing they are bounded like anything else. assert effective_packing({**text, **raw}, branch_never_packs = True) is False assert max_train_rows_for_config({**text, **raw}, branch_never_packs = True) == 1024 @@ -897,7 +915,9 @@ def test_both_loaders_apply_the_row_bound(): import ast from pathlib import Path - worker_src = (Path(__file__).resolve().parents[1] / "core/training/worker.py").read_text() + worker_src = (Path(__file__).resolve().parents[1] / "core/training/worker.py").read_text( + encoding = "utf-8" + ) tree = ast.parse(worker_src) calls = {} for node in ast.walk(tree): From 8c9133899dc1f28acdd959870461f883ff7bee08 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 15 Aug 2026 16:46:56 +0000 Subject: [PATCH 13/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_training_preflight.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/studio/backend/tests/test_training_preflight.py b/studio/backend/tests/test_training_preflight.py index f79be6b47..6653a583a 100644 --- a/studio/backend/tests/test_training_preflight.py +++ b/studio/backend/tests/test_training_preflight.py @@ -649,9 +649,7 @@ def test_effective_packing_decides_the_opt_out(): # `not raw_text_mode` and gives way to the text path, while audio # preprocessing is chosen before the raw-text bypass and holds either way. for raw in ({"training_type": "Continued Pretraining"}, {"format_type": "raw"}): - assert ( - effective_packing({**text, **raw, "packing": True}, branch_never_packs = True) is False - ) + assert effective_packing({**text, **raw, "packing": True}, branch_never_packs = True) is False assert effective_packing({**text, **raw, "packing": True}) is True # Without packing they are bounded like anything else. assert effective_packing({**text, **raw}, branch_never_packs = True) is False From 170a6e717e8204d24521756a5b4e7f4db2f8fa38 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 15 Aug 2026 16:47:15 +0000 Subject: [PATCH 14/14] Add staging CI workflows for unslothai/unsloth#8890 --- .github/workflows/staging-8890-macos-14.yml | 30 +++++++++++++++++++ .../workflows/staging-8890-ubuntu-latest.yml | 30 +++++++++++++++++++ .../workflows/staging-8890-windows-latest.yml | 30 +++++++++++++++++++ 3 files changed, 90 insertions(+) create mode 100644 .github/workflows/staging-8890-macos-14.yml create mode 100644 .github/workflows/staging-8890-ubuntu-latest.yml create mode 100644 .github/workflows/staging-8890-windows-latest.yml diff --git a/.github/workflows/staging-8890-macos-14.yml b/.github/workflows/staging-8890-macos-14.yml new file mode 100644 index 000000000..57accd6ee --- /dev/null +++ b/.github/workflows/staging-8890-macos-14.yml @@ -0,0 +1,30 @@ +name: "staging-8890 macos-14" +on: + push: + branches: ["pr-8890-xplat-ci"] + workflow_dispatch: +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true +permissions: + contents: read +defaults: + run: + shell: bash +jobs: + test: + runs-on: macos-14 + timeout-minutes: 30 + env: + UNSLOTH_COMPILE_DISABLE: '1' + UNSLOTH_IS_PRESENT: '1' + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - uses: actions/setup-python@v5 + with: + python-version: '3.12' + - run: python -m pip install -r studio/backend/requirements/studio.txt pytest pytest-asyncio pytest-timeout python-multipart + - run: python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + - run: rc=0; PYTHONPATH=studio/backend python -m pytest studio/backend/tests/test_training_preflight.py -q --tb=short -k 'not llama_cpp_load_progress_live and not TestGpuAutoSelection and not TestPreSpawnGpuResolution and not TestPerGpuFitGuardAllCounts and not TestTransformersIntrospection and not test_returns_cuda_when_cuda_available and not test_calls_cuda_cache_when_cuda' > pytest_out.txt 2>&1 || rc=$?; cat pytest_out.txt; if [ "$rc" = "5" ]; then echo "no tests ran (deps absent on this runner)"; rc=0; fi; if [ "$rc" = "2" ] && grep -qE "No module named .(torch|unsloth_zoo|transformers)." pytest_out.txt && ! grep -qE "^(FAILED|ERROR) " pytest_out.txt; then echo "collection needs a dep this runner does not ship"; rc=0; fi; if [ "$rc" = "1" ] && ! grep -qE "^ERROR " pytest_out.txt; then tot=$(grep -cE "^FAILED " pytest_out.txt); dep=$(grep -cE "^FAILED .* No module named .(torch|unsloth_zoo|transformers).$" pytest_out.txt); if [ "$tot" -gt 0 ] && [ "$tot" = "$dep" ]; then echo "only tests needing a dep this runner does not ship failed"; rc=0; fi; fi; exit "$rc" diff --git a/.github/workflows/staging-8890-ubuntu-latest.yml b/.github/workflows/staging-8890-ubuntu-latest.yml new file mode 100644 index 000000000..8f4510ccc --- /dev/null +++ b/.github/workflows/staging-8890-ubuntu-latest.yml @@ -0,0 +1,30 @@ +name: "staging-8890 ubuntu-latest" +on: + push: + branches: ["pr-8890-xplat-ci"] + workflow_dispatch: +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true +permissions: + contents: read +defaults: + run: + shell: bash +jobs: + test: + runs-on: ubuntu-latest + timeout-minutes: 30 + env: + UNSLOTH_COMPILE_DISABLE: '1' + UNSLOTH_IS_PRESENT: '1' + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - uses: actions/setup-python@v5 + with: + python-version: '3.12' + - run: python -m pip install -r studio/backend/requirements/studio.txt pytest pytest-asyncio pytest-timeout python-multipart + - run: python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + - run: rc=0; PYTHONPATH=studio/backend python -m pytest studio/backend/tests/test_training_preflight.py -q --tb=short -k 'not llama_cpp_load_progress_live and not TestGpuAutoSelection and not TestPreSpawnGpuResolution and not TestPerGpuFitGuardAllCounts and not TestTransformersIntrospection and not test_returns_cuda_when_cuda_available and not test_calls_cuda_cache_when_cuda' > pytest_out.txt 2>&1 || rc=$?; cat pytest_out.txt; if [ "$rc" = "5" ]; then echo "no tests ran (deps absent on this runner)"; rc=0; fi; if [ "$rc" = "2" ] && grep -qE "No module named .(torch|unsloth_zoo|transformers)." pytest_out.txt && ! grep -qE "^(FAILED|ERROR) " pytest_out.txt; then echo "collection needs a dep this runner does not ship"; rc=0; fi; if [ "$rc" = "1" ] && ! grep -qE "^ERROR " pytest_out.txt; then tot=$(grep -cE "^FAILED " pytest_out.txt); dep=$(grep -cE "^FAILED .* No module named .(torch|unsloth_zoo|transformers).$" pytest_out.txt); if [ "$tot" -gt 0 ] && [ "$tot" = "$dep" ]; then echo "only tests needing a dep this runner does not ship failed"; rc=0; fi; fi; exit "$rc" diff --git a/.github/workflows/staging-8890-windows-latest.yml b/.github/workflows/staging-8890-windows-latest.yml new file mode 100644 index 000000000..9d80daa67 --- /dev/null +++ b/.github/workflows/staging-8890-windows-latest.yml @@ -0,0 +1,30 @@ +name: "staging-8890 windows-latest" +on: + push: + branches: ["pr-8890-xplat-ci"] + workflow_dispatch: +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true +permissions: + contents: read +defaults: + run: + shell: bash +jobs: + test: + runs-on: windows-latest + timeout-minutes: 30 + env: + UNSLOTH_COMPILE_DISABLE: '1' + UNSLOTH_IS_PRESENT: '1' + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - uses: actions/setup-python@v5 + with: + python-version: '3.12' + - run: python -m pip install -r studio/backend/requirements/studio.txt pytest pytest-asyncio pytest-timeout python-multipart + - run: python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + - run: rc=0; PYTHONPATH=studio/backend python -m pytest studio/backend/tests/test_training_preflight.py -q --tb=short -k 'not llama_cpp_load_progress_live and not TestGpuAutoSelection and not TestPreSpawnGpuResolution and not TestPerGpuFitGuardAllCounts and not TestTransformersIntrospection and not test_returns_cuda_when_cuda_available and not test_calls_cuda_cache_when_cuda' > pytest_out.txt 2>&1 || rc=$?; cat pytest_out.txt; if [ "$rc" = "5" ]; then echo "no tests ran (deps absent on this runner)"; rc=0; fi; if [ "$rc" = "2" ] && grep -qE "No module named .(torch|unsloth_zoo|transformers)." pytest_out.txt && ! grep -qE "^(FAILED|ERROR) " pytest_out.txt; then echo "collection needs a dep this runner does not ship"; rc=0; fi; if [ "$rc" = "1" ] && ! grep -qE "^ERROR " pytest_out.txt; then tot=$(grep -cE "^FAILED " pytest_out.txt); dep=$(grep -cE "^FAILED .* No module named .(torch|unsloth_zoo|transformers).$" pytest_out.txt); if [ "$tot" -gt 0 ] && [ "$tot" = "$dep" ]; then echo "only tests needing a dep this runner does not ship failed"; rc=0; fi; fi; exit "$rc"