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" diff --git a/studio/backend/core/training/dataset_bounds.py b/studio/backend/core/training/dataset_bounds.py new file mode 100644 index 000000000..e514939d7 --- /dev/null +++ b/studio/backend/core/training/dataset_bounds.py @@ -0,0 +1,265 @@ +# 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 +import re +import tempfile +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 +# 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: + """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, 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 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 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. Pass the + 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 + return not branch_never_packs + + +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, branch_never_packs = branch_never_packs): + 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 run_dir_for_checkpoint(checkpoint_path: Any) -> Optional[str]: + """The run directory a checkpoint lives in, or None when there is none. + + 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: + return None + path = str(checkpoint_path).rstrip("/\\") + if not path: + return None + head, tail = os.path.split(path) + 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 + + +def record_row_bound( + output_dir: Any, + max_train_rows: Optional[int], + seed: Any = 3407, +) -> 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 + 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. + + 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 + -- 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 False + marker = os.path.join(run_dir, ROW_BOUND_MARKER_FILE) + tmp_path = None + try: + 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 False + finally: + if tmp_path is not None: + try: + os.unlink(tmp_path) + except OSError: + pass + return True + + +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(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 None, fallback_seed + return _positive_int(recorded, 0) or None, _seed_int(marker.get("seed"), fallback_seed) + + +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 7f4d9471c..67c63405b 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -89,6 +89,8 @@ should_use_mlx_training_backend, ) +from .dataset_bounds import bound_dataset_rows + logger = get_logger(__name__) # Streaming eval has no __len__, so an unbounded eval would iterate the whole @@ -2572,6 +2574,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 +2583,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 +3024,36 @@ 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. + # 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): + logger.info( + 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 {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 8db7e6462..3a62fbbcd 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -797,7 +797,14 @@ 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 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, diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index d4f7be7e9..068eeef86 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -54,6 +54,15 @@ 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, + 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 ( direct_wheel_url, @@ -2779,14 +2788,48 @@ 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. + # 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 + # 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 + ) + + # 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 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)))) + if mlx_split_names_rows: + return ds + return bound_dataset_rows( + ds, + mlx_max_train_rows, + mlx_max_train_rows_seed, + 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 @@ -2993,6 +3036,17 @@ 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. + 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) @@ -4084,6 +4138,11 @@ def _apply_stop(save: bool) -> None: training_type = config.get("training_type", "LoRA/QLoRA") is_cpt_for_dataset = training_type == "Continued Pretraining" + # 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) + 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 +4166,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 = max_train_rows_seed, ) if isinstance(result, tuple): loaded_dataset, loaded_eval_dataset = result @@ -4165,6 +4226,35 @@ 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. + # 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 + # 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() @@ -4463,6 +4553,17 @@ 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. + 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 0c99fc3cd..6653a583a 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,503 @@ 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.dataset_bounds 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_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) + + 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_max_steps_dataset_rows_survives_unusable_numbers(): + 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. + 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 + + # 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}, 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 + # "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 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 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 + + +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_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) + + 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) + + # 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 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_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 + + # 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 + + # 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) + + # 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 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(): + 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( + encoding = "utf-8" + ) + 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"] + # 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] + + # 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 + + 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)',