diff --git a/README.md b/README.md index f1405f8..616075e 100644 --- a/README.md +++ b/README.md @@ -10,6 +10,7 @@ This repo supports two training backends: - [Quick Start](#quick-start) - [Run SFT using a pre-built container](#-run-sft-using-a-pre-built-container) +- [Reasoning SFT on LUMI](docs/reasoning-sft.md) - [Project Structure](#-project-structure) - [Design Philosophy](#-design-philosophy) - [Feature Guide](#-feature-guide) @@ -116,6 +117,12 @@ Tokenizing first keeps the multi-node allocation from sitting idle during CPU-bo The example config is [`configs/trl/prelude-sft.yaml`](configs/trl/prelude-sft.yaml). It fine-tunes a 9B checkpoint on LUMI, with the tokenizer from a separate repo. For your own run, copy it and replace the checkpoint, data, container paths, and SLURM account. +For reasoning continuation, use the production-oriented +[`configs/trl/reasoning-sft-lumi.yaml`](configs/trl/reasoning-sft-lumi.yaml) and +the accompanying [reasoning SFT guide](docs/reasoning-sft.md). It adds a frozen +FSDP launch profile, a pinned parent revision, 16K assistant-only training, and +an explicit contract for a token-balanced materialized reasoning mixture. + ### Before you start The login node needs only the base dependencies from [Installation](#installation), because the training stack lives in the container. Run every `submit.py` command from the repository root, inside that environment. Relative paths in the config (`container.env_file`, `paths.output_base`) resolve against the root, and `submit.py` copies the code from it. @@ -305,6 +312,7 @@ data: datasets: - name: "my_dataset" path: "org/dataset" + revision: "" # pin Hub inputs for reproducible runs split: "train" weight: 1.0 # 1 = full dataset, <1 undersamples, >1 oversamples ``` @@ -399,7 +407,7 @@ You must specify exactly one determining factor for training duration in the `tr - **DeepSpeed**: configured inline under the top-level `deepspeed:` key (see the reference config below); set `deepspeed: null` to disable DeepSpeed entirely. To switch from ZeRO stage 2 to stage 3, bump `zero_optimization.stage` to `3` and add the `stage3_*` tuning keys — see `configs/deepspeed/zero3.yaml` for a full example. -- **Accelerate flags**: the `accelerate` section in the YAML mirrors the CLI flags required for multi-node setups (`mixed_precision`, `dynamo_backend`, `rdzv_backend`, etc.). +- **Accelerate flags**: the `accelerate` section in the YAML mirrors the CLI flags required for multi-node setups (`mixed_precision`, `dynamo_backend`, `rdzv_backend`, etc.). Set `accelerate.config_file` for a full profile such as FSDP. The renderer copies that file into the run's `slurm/` directory before submission, so a queued job cannot drift with the working checkout. These are used by the SLURM launcher to generate the correct job script. - **Self-healing**: the SLURM launcher (`src/post_training/slurm/`) supports auto-requeueing. - `slurm.signal_time_seconds` ensures the job saves a checkpoint and requeues itself before the wall time expires diff --git a/configs/accelerate/fsdp-lumi.yaml b/configs/accelerate/fsdp-lumi.yaml new file mode 100644 index 0000000..8af88f1 --- /dev/null +++ b/configs/accelerate/fsdp-lumi.yaml @@ -0,0 +1,14 @@ +# Multi-node full-shard profile for LUMI MI250X GCDs. The generated SLURM +# launcher overrides process counts, machine rank, and rendezvous coordinates. +compute_environment: LOCAL_MACHINE +distributed_type: FSDP +mixed_precision: bf16 +num_processes: 8 +fsdp_config: + fsdp_sharding_strategy: FULL_SHARD + fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP + fsdp_state_dict_type: FULL_STATE_DICT + fsdp_use_orig_params: true + fsdp_offload_params: false + fsdp_cpu_ram_efficient_loading: true + fsdp_sync_module_states: true diff --git a/configs/trl/prelude-sft.yaml b/configs/trl/prelude-sft.yaml index 33e5b2a..bcdc1b1 100644 --- a/configs/trl/prelude-sft.yaml +++ b/configs/trl/prelude-sft.yaml @@ -54,6 +54,7 @@ data: datasets: - name: "dolci-instruct-sft" path: "allenai/Dolci-Instruct-SFT" # HuggingFace dataset path + revision: bd3c8f3a9b2cc5a9682e44b96ddd0bb2ff027221 split: "train" weight: 1.0 transform: null # null = already conversational diff --git a/configs/trl/reasoning-sft-lumi.yaml b/configs/trl/reasoning-sft-lumi.yaml new file mode 100644 index 0000000..b96909a --- /dev/null +++ b/configs/trl/reasoning-sft-lumi.yaml @@ -0,0 +1,117 @@ +# Reasoning continuation profile validated against the OpenEuroLLM 9B Qwen3 +# checkpoint family. Replace only the site-specific paths and the materialized +# data artifact before running the tokenize-only gate described in the guide. +method: sft +backend: trl +run_name: oellm-9b-anneal300b-reasoning-sft +offline: true +prefetch_assets: true + +container: + image: /scratch//containers/post-training-rocm7.2.4-py3.12-torch2.9.1-trl1.7.0-olmo-patched.sif + bind_mounts: + - /scratch/ + path: /opt/venv/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin + env_file: env/lumi.env + +model: + name_or_path: Neonkraft/oellm-9b-256k-theta64m-prelude-anneal300b-instruct-sft + revision: 85bf18fb4f0bee6ac6270f06b1d1c6b3be200f31 + attn_implementation: flash_attention_2 + dtype: bfloat16 + tokenizer_name_or_path: null + tokenizer_revision: null + +training: + # 64 packed 16K sequences/step * 500 steps = 524,288,000 training tokens. + num_training_tokens: 524288000 + learning_rate: 1.5e-6 + effective_batch_size: 64 + per_device_train_batch_size: 1 + warmup_steps: 15 + lr_scheduler_type: cosine + lr_scheduler_kwargs: null + adam_beta1: 0.9 + adam_beta2: 0.95 + weight_decay: 0.0 + adam_epsilon: 1.0e-8 + gradient_checkpointing: true + gradient_checkpointing_kwargs: + use_reentrant: false + bf16: true + seed: 20260917 + use_liger_kernel: true + liger_kernel_config: + fused_linear_cross_entropy: true + cross_entropy: false + rope: false + rms_norm: false + swiglu: false + +sft: + max_seq_length: 16384 + packing: true + dataset_num_proc: 6 + # Never train a reasoning trace cut before or its final answer. + truncated_span_action: drop + +checkpointing: + save_strategy: steps + save_steps: 100 + save_total_limit: 5 + inference_checkpoint_steps: 100 + inference_checkpoint_path: inference_checkpoints + +data: + chat_template: qwen3 + num_proc: 6 + seed: 20260917 + datasets: + # The production recipe is token-balanced, globally deduplicated and quality + # gated before training. Point this at that immutable, materialized artifact; + # do not replace it with row weights over the raw source datasets. + - name: reasoning-anneal300b-v1 + path: /scratch///post-training-data/reasoning-anneal300b-v1/train.parquet + split: train + weight: 1.0 + transform: null + +# FSDP is the tested 8-node path for this 9B/16K profile. Keep DeepSpeed off so +# Accelerate uses the frozen profile below. +deepspeed: null + +accelerate: + config_file: configs/accelerate/fsdp-lumi.yaml + mixed_precision: bf16 + use_deepspeed: false + same_network: true + rdzv_backend: static + dynamo_backend: "no" + +logging: + report_to: [] + wandb_project: oellm-reasoning + logging_steps: 10 + include_num_input_tokens_seen: non_padding + +slurm: + account: + partition: standard-g + num_nodes: 8 + gpus_per_node: 8 + cpus_per_task: 56 + mem: 256G + wall_time: "24:00:00" + job_name: oellm-reasoning-sft + signal_time_seconds: 300 + max_failures: 1 + modules: [] + module_purge: false + +debug: + enabled: false + override_existing: false + +paths: + output_base: outputs + debug_base: outputs/debug diff --git a/configs/trl/sft.yaml b/configs/trl/sft.yaml index c42689d..11b2263 100644 --- a/configs/trl/sft.yaml +++ b/configs/trl/sft.yaml @@ -89,6 +89,7 @@ data: datasets: - name: "nemotron_pt_v2" path: "nvidia/Nemotron-Post-Training-Dataset-v2" + revision: null # pin production runs to a Hub commit split: "stem" weight: 1.0 # 1 = full dataset, <1 undersamples, >1 oversamples transform: null # null = already conversational @@ -156,4 +157,3 @@ paths: output_base: "outputs" debug_base: "outputs/debug" - diff --git a/docs/reasoning-sft.md b/docs/reasoning-sft.md new file mode 100644 index 0000000..330ad42 --- /dev/null +++ b/docs/reasoning-sft.md @@ -0,0 +1,121 @@ +# Reasoning SFT on LUMI + +Reasoning continuation is supervised fine-tuning: the trainer does not need a +separate `reasoning` method. It does need stricter data preparation and launch +discipline than a general instruction run. The reference profile is +[`configs/trl/reasoning-sft-lumi.yaml`](../configs/trl/reasoning-sft-lumi.yaml). + +## What the framework already enforces + +The TRL path is suitable for a production reasoning run because it: + +- applies the Qwen3 chat template with generation markers and refuses to start + if assistant-only loss cannot be constructed; +- drops rows whose assistant answer contributes no supervised tokens; +- can drop a row when the 16K boundary cuts through its supervised span, which + prevents training on an unterminated `` block or final answer; +- packs complete conversational records and derives the step count from a + training-token budget; +- aligns generation EOS with the chat template's `<|im_end|>` turn terminator; +- pins the parent model, tokenizer, and Hub datasets to revisions; +- freezes the run config, source, Accelerate profile, and generated SLURM job; +- supports a one-GPU `--tokenize-only` gate before the full allocation; and +- resumes full checkpoints and writes lightweight inference checkpoints. + +The reference profile adds the multi-node FSDP configuration used by the 9B, +16K LUMI run. The generated job still supplies the live node count, process +count, machine rank and rendezvous address, so the same profile works for the +one-GPU gate and the eight-node run. + +## Data contract + +The training input is one immutable, materialized Parquet dataset with a +`messages` column. Each value is a list of `{role, content}` messages and the +last turn is the assistant answer. Reasoning examples contain a non-empty +`...` block followed by a non-empty final answer. Replay examples +may contain a normal answer without a think block. + +Prepare the artifact before this trainer sees it. A production builder should: + +1. pin every source to a commit revision; +2. normalize records to the `messages` contract without flattening prompt and + answer into one string; +3. keep only complete reasoning traces and final answers; +4. reject overlength rows rather than truncate them; +5. reject exact lexical loops (the current run uses a repeated 30-gram at least + 20 times as the strict threshold); +6. deduplicate normalized user prompts across sources, in a deterministic + priority order; +7. allocate mixture shares by rendered tokens, not rows; and +8. write a manifest with source revisions, filter counts, selected row/token + counts, tokenizer/template identity, seed, and output SHA-256. + +The framework's ordinary `data.datasets[].weight` is deliberately a row-count +multiplier. It is useful for similarly distributed instruction datasets, but +it does **not** reproduce a token-balanced reasoning mixture whose traces range +from hundreds to tens of thousands of tokens. The reference config therefore +consumes the already-materialized mixture as one dataset. + +For the current conservative continuation recipe, the materialized artifact is +524,288,000 rendered tokens: a 65% reasoning / 35% instruction-replay target, +plus a fixed multilingual reasoning coverage floor consumed once before the +weighted allocation. The exact source recipe and its manifest belong with the +data artifact; the training repository records the immutable artifact used by +the run. + +## Configure the run + +Copy the reference YAML and replace these site-specific values: + +- `container.image` and `container.bind_mounts`; +- `container.env_file` with the shared Hugging Face cache variables; +- `data.datasets[0].path` with the immutable Parquet artifact; +- `slurm.account`; and +- `paths.output_base` if the default is not on a persistent filesystem. + +Do not change the pinned model revision, Qwen3 template, 16K sequence length, +assistant-only loss behavior, or Liger fused-loss settings without a new smoke +test. The parent checkpoint contains its tokenizer, so the tokenizer follows +the same pinned revision. + +## Gate, train, verify + +First render and tokenize through the same code path as production: + +```bash +python scripts/submit.py \ + --config configs/trl/reasoning-sft-lumi.yaml \ + --tokenize-only +``` + +Read the generated `slurm-.out` and `.err`. The gate passes only when: + +- the decoded preview uses Qwen3 ChatML and has a supervised assistant answer; +- no row is cut part-way through its supervised span; +- the resolved step count is 500; +- the parent and dataset paths resolve from the shared cache/filesystem; and +- trainer construction completes without an FSDP, FlashAttention or Liger + error. + +Then submit the unchanged configuration: + +```bash +python scripts/submit.py --config configs/trl/reasoning-sft-lumi.yaml +``` + +After training starts, confirm that the full run reuses the tokenization cache, +that loss is finite and decreasing, that every rank advances together, and +that a resumable checkpoint is written at step 100. A release candidate should +be compared with the parent on the same multilingual instruction, math/code +reasoning, long-context, and repetition-loop probes. Publish the parent commit, +data-manifest checksum, frozen config, environment, SLURM job IDs, selected +checkpoint, and evaluation results in the model card. + +## Scope still outside this repository + +This trainer does not currently build token-balanced, cross-source-deduplicated +mixtures, perform benchmark decontamination, or run the evaluation suite. Those +steps should stay fail-closed in the data/evaluation pipelines and hand this +repository immutable artifacts plus manifests. Keeping that boundary explicit +prevents a convenient row-weighted training config from being mistaken for the +audited production data recipe. diff --git a/scripts/data.py b/scripts/data.py index 64c3a53..419f89b 100755 --- a/scripts/data.py +++ b/scripts/data.py @@ -128,6 +128,8 @@ def _run_inspect(args: argparse.Namespace, cli_overrides: list[str]) -> None: load_kwargs["name"] = entry.subset if getattr(entry, "data_dir", None) is not None: load_kwargs["data_dir"] = entry.data_dir + if getattr(entry, "revision", None) is not None and not Path(entry.path).exists(): + load_kwargs["revision"] = entry.revision ds = load_dataset(entry.path, split=entry.split, **load_kwargs) diff --git a/src/post_training/backend.py b/src/post_training/backend.py index b3f4685..0451013 100644 --- a/src/post_training/backend.py +++ b/src/post_training/backend.py @@ -71,11 +71,26 @@ def _shorten_dataset_name(name: str) -> str: def _dataset_mix_hash(datasets: list) -> str: + # Every field that can change the realised rows belongs in the identity. + # Omitting (for example) the Hub revision can make two materially different + # runs resolve to the same output directory and resume each other's state. canonical = sorted( - [{"name": d.name, "path": d.path, "weight": d.weight} for d in datasets], + [ + { + "name": d.name, + "path": d.path, + "revision": d.revision, + "data_dir": d.data_dir, + "subset": d.subset, + "split": d.split, + "weight": d.weight, + "transform": d.transform, + } + for d in datasets + ], key=lambda x: x["name"], ) - return hashlib.sha256(json.dumps(canonical).encode()).hexdigest()[:8] + return hashlib.sha256(json.dumps(canonical, sort_keys=True).encode()).hexdigest()[:8] class TRLBackend(Backend): @@ -168,7 +183,7 @@ def generate_run_name(self, config: PostTrainingConfig, timestamp: str) -> str: model_short = _shorten_model_name(config.model.name_or_path) datasets = config.data.datasets if len(datasets) == 1: - ds_part = _shorten_dataset_name(datasets[0].name) + ds_part = f"{_shorten_dataset_name(datasets[0].name)}_{_dataset_mix_hash(datasets)}" else: ds_part = f"mix_{_dataset_mix_hash(datasets)}" return f"{config.method}-{model_short}-{ds_part}-{timestamp}" diff --git a/src/post_training/config.py b/src/post_training/config.py index 6a20f1a..f6c15a4 100644 --- a/src/post_training/config.py +++ b/src/post_training/config.py @@ -178,6 +178,8 @@ class DatasetEntry: name: str = MISSING path: str = MISSING + # Hub branch, tag, or commit. Ignored for existing local paths. + revision: str | None = None data_dir: str | None = None subset: str | None = None split: str = "train" @@ -199,6 +201,10 @@ class DataConfig: class AccelerateConfig: """Flags forwarded to ``accelerate launch`` for explicit multi-node control.""" + # Optional Accelerate YAML for launchers that are not expressible by the + # simple flags below (for example FSDP on LUMI). The SLURM renderer freezes + # a copy beside job.sh before submission. + config_file: str | None = None mixed_precision: str = "bf16" # Only takes effect when the top-level `deepspeed:` config is also set; # `deepspeed: null` disables DeepSpeed at launch regardless of this flag. diff --git a/src/post_training/data/loader.py b/src/post_training/data/loader.py index 4d38677..c1599b6 100644 --- a/src/post_training/data/loader.py +++ b/src/post_training/data/loader.py @@ -153,10 +153,11 @@ def load_and_mix_datasets( loaded_datasets: list[Dataset] = [] for entry in entries: logger.info( - "Loading dataset '%s' from '%s' (data_dir=%s, subset=%s, split=%s, " - "weight=%s, transform=%s)", + "Loading dataset '%s' from '%s' (revision=%s, data_dir=%s, subset=%s, " + "split=%s, weight=%s, transform=%s)", entry.name, entry.path, + entry.revision, entry.data_dir, entry.subset, entry.split, @@ -170,6 +171,8 @@ def load_and_mix_datasets( load_kwargs["data_dir"] = entry.data_dir if entry.subset is not None: load_kwargs["name"] = entry.subset + if entry.revision is not None and not Path(entry.path).exists(): + load_kwargs["revision"] = entry.revision ds = _load_dataset_entry(entry.path, entry.split, **load_kwargs) diff --git a/src/post_training/slurm/job.sh.jinja b/src/post_training/slurm/job.sh.jinja index 8cf99d7..8bf0a1f 100644 --- a/src/post_training/slurm/job.sh.jinja +++ b/src/post_training/slurm/job.sh.jinja @@ -61,6 +61,8 @@ trap handle_signal USR1 # ── Launch via accelerate (explicit flags from config) ───────────────────── export LAUNCHER="accelerate launch \ +{% if accelerate_config_file %} --config_file {{ accelerate_config_file }} \ +{% endif -%} --num_machines $NUM_NODES \ --num_processes $WORLD_SIZE \ --machine_rank \$SLURM_PROCID \ @@ -96,4 +98,3 @@ if [ $EXIT_CODE -ne 0 ]; then exit 1 fi fi - diff --git a/src/post_training/slurm/job_trl_container.sh.jinja b/src/post_training/slurm/job_trl_container.sh.jinja index 989a947..3cd2adb 100644 --- a/src/post_training/slurm/job_trl_container.sh.jinja +++ b/src/post_training/slurm/job_trl_container.sh.jinja @@ -110,6 +110,8 @@ srun --export=ALL --wait=60 --kill-on-bad-exit=1 \ export WANDB_DIR=\"{{ run_dir }}\" accelerate launch \ +{% if accelerate_config_file %} --config_file {{ accelerate_config_file }} \ +{% endif -%} --num_machines \$NNODES \ --num_processes $WORLD_SIZE \ --machine_rank \$NODE_RANK \ diff --git a/src/post_training/slurm/launcher.py b/src/post_training/slurm/launcher.py index fd8338f..5340578 100644 --- a/src/post_training/slurm/launcher.py +++ b/src/post_training/slurm/launcher.py @@ -8,6 +8,7 @@ from __future__ import annotations import logging +import shutil import subprocess from pathlib import Path from typing import TYPE_CHECKING @@ -25,6 +26,28 @@ _TRL_CONTAINER_TEMPLATE_NAME = "job_trl_container.sh.jinja" +def _freeze_accelerate_config(config: PostTrainingConfig, run_dir: Path) -> str | None: + """Freeze the optional Accelerate launch profile next to ``job.sh``. + + A queued job must not read a mutable file from the working checkout. The + copied path is also inside ``run_dir``, which the container launcher already + exposes to the training process. + """ + source_value = config.accelerate.config_file + if source_value is None: + return None + + source = Path(source_value).expanduser().resolve() + if not source.is_file(): + raise FileNotFoundError(f"accelerate.config_file '{source_value}' not found.") + + destination = run_dir / "slurm" / "accelerate.yaml" + destination.parent.mkdir(parents=True, exist_ok=True) + if source != destination.resolve(): + shutil.copy2(source, destination) + return str(destination.resolve()) + + def render_trl_slurm_script( config: PostTrainingConfig, run_dir: Path, @@ -54,6 +77,7 @@ def render_trl_slurm_script( keep_trailing_newline=True, ) template = env.get_template(_TEMPLATE_NAME) + accelerate_config_file = _freeze_accelerate_config(config, run_dir) rendered = template.render( # SLURM parameters @@ -75,6 +99,7 @@ def render_trl_slurm_script( config_path=config_path, tokenize_only=tokenize_only, # Accelerate flags + accelerate_config_file=accelerate_config_file, mixed_precision=config.accelerate.mixed_precision, dynamo_backend=config.accelerate.dynamo_backend, use_deepspeed=config.accelerate.use_deepspeed and bool(config.deepspeed), @@ -113,6 +138,7 @@ def render_trl_container_slurm_script( keep_trailing_newline=True, ) template = env.get_template(_TRL_CONTAINER_TEMPLATE_NAME) + accelerate_config_file = _freeze_accelerate_config(config, run_dir) rendered = template.render( # SLURM parameters @@ -134,6 +160,7 @@ def render_trl_container_slurm_script( config_path=config_path, tokenize_only=tokenize_only, # Accelerate flags + accelerate_config_file=accelerate_config_file, mixed_precision=config.accelerate.mixed_precision, dynamo_backend=config.accelerate.dynamo_backend, use_deepspeed=config.accelerate.use_deepspeed and bool(config.deepspeed), diff --git a/src/post_training/utils/prefetch.py b/src/post_training/utils/prefetch.py index 3097114..d057e68 100644 --- a/src/post_training/utils/prefetch.py +++ b/src/post_training/utils/prefetch.py @@ -113,12 +113,19 @@ def _prefetch_dataset(entry: DatasetEntry) -> None: if _is_local(entry.path): logger.info("Dataset '%s' is a local path, skipping download.", entry.name) return - logger.info("Downloading dataset '%s' ('%s') to HF cache...", entry.name, entry.path) + logger.info( + "Downloading dataset '%s' ('%s', revision=%s) to HF cache...", + entry.name, + entry.path, + entry.revision or "main", + ) load_kwargs: dict = {} if entry.data_dir is not None: load_kwargs["data_dir"] = entry.data_dir if entry.subset is not None: load_kwargs["name"] = entry.subset + if entry.revision is not None: + load_kwargs["revision"] = entry.revision load_dataset(entry.path, split=entry.split, **load_kwargs) logger.info("Dataset '%s' cached.", entry.name) diff --git a/tests/test_backend.py b/tests/test_backend.py index eebe4d3..677e83d 100644 --- a/tests/test_backend.py +++ b/tests/test_backend.py @@ -2,8 +2,35 @@ import logging -from post_training.backend import TRLBackend -from post_training.config import ContainerConfig, PostTrainingConfig +from post_training.backend import TRLBackend, _dataset_mix_hash +from post_training.config import ContainerConfig, DatasetEntry, PostTrainingConfig + + +def test_dataset_mix_hash_covers_reproducibility_fields(): + baseline = DatasetEntry(name="reasoning", path="org/data") + changes = [ + DatasetEntry(name="reasoning", path="org/data", revision="abc123"), + DatasetEntry(name="reasoning", path="org/data", subset="math"), + DatasetEntry(name="reasoning", path="org/data", split="validation"), + DatasetEntry(name="reasoning", path="org/data", transform="custom"), + ] + + baseline_hash = _dataset_mix_hash([baseline]) + + assert all(_dataset_mix_hash([changed]) != baseline_hash for changed in changes) + + +def test_single_dataset_run_name_changes_with_revision(): + config = PostTrainingConfig() + config.model.name_or_path = "org/model" + config.data.datasets = [DatasetEntry(name="reasoning", path="org/data")] + first = TRLBackend().generate_run_name(config, "20260917-120000") + + config.data.datasets[0].revision = "abc123" + second = TRLBackend().generate_run_name(config, "20260917-120000") + + assert first != second + assert first.startswith("sft-model-reasoning_") def _make_fake_repo(repo_dir, marker: str = "live") -> None: diff --git a/tests/test_config.py b/tests/test_config.py index ee30cbc..dbe69a4 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,5 +1,7 @@ """Tests for nullable config fields.""" +from pathlib import Path + import pytest import yaml @@ -203,3 +205,16 @@ def test_resolve_tokenizer(tokenizer_name_or_path, tokenizer_revision, expected) def test_resolve_tokenizer_without_any_revision(): assert ModelConfig(name_or_path="org/model").resolve_tokenizer() == ("org/model", None) + + +def test_reasoning_lumi_profile_resolves_exact_production_budget(): + config_path = Path(__file__).parents[1] / "configs" / "trl" / "reasoning-sft-lumi.yaml" + + config = PostTrainingConfig.load(config_path) + + assert config.resolve_max_steps() == 500 + assert config.model.revision == "85bf18fb4f0bee6ac6270f06b1d1c6b3be200f31" + assert config.data.chat_template == "qwen3" + assert config.sft.truncated_span_action == "drop" + assert config.accelerate.config_file == "configs/accelerate/fsdp-lumi.yaml" + assert config.deepspeed is None diff --git a/tests/test_data_loader.py b/tests/test_data_loader.py index 8cabfc6..155fa98 100644 --- a/tests/test_data_loader.py +++ b/tests/test_data_loader.py @@ -71,6 +71,42 @@ def test_weight_resampling_uses_each_dataset_size(monkeypatch): assert Counter(mixed["source"]) == {"a": 2, "b": 4} +def test_hub_dataset_revision_is_forwarded(monkeypatch): + calls = [] + + def fake_load_dataset(path: str, **kwargs) -> Dataset: + calls.append((path, kwargs)) + return _dataset("a", 1) + + monkeypatch.setattr(loader, "load_dataset", fake_load_dataset) + + loader.load_and_mix_datasets( + _config(DatasetEntry(name="a", path="org/dataset", revision="abc123")) + ) + + assert calls == [("org/dataset", {"split": "train", "revision": "abc123"})] + + +def test_local_dataset_file_does_not_receive_hub_revision(monkeypatch, tmp_path): + path = tmp_path / "train.parquet" + path.touch() + calls = [] + + def fake_load_dataset(builder: str, **kwargs) -> Dataset: + calls.append((builder, kwargs)) + return _dataset("a", 1) + + monkeypatch.setattr(loader, "load_dataset", fake_load_dataset) + + loader.load_and_mix_datasets( + _config(DatasetEntry(name="a", path=str(path), revision="ignored-for-local")) + ) + + assert calls == [ + ("parquet", {"data_files": str(path), "split": "train"}), + ] + + def test_weight_resampling_is_reproducible(monkeypatch): _patch_load_dataset( monkeypatch, diff --git a/tests/test_prefetch.py b/tests/test_prefetch.py index a957c63..212eda7 100644 --- a/tests/test_prefetch.py +++ b/tests/test_prefetch.py @@ -2,7 +2,7 @@ import pytest -from post_training.config import PostTrainingConfig +from post_training.config import DatasetEntry, PostTrainingConfig from post_training.utils import prefetch @@ -78,3 +78,29 @@ def test_local_tokenizer_path_is_not_downloaded(downloads, tmp_path): assert downloads == [("org/model", None, None)] assert paths.tokenizer == str(tmp_path) + + +def test_dataset_revision_is_prefetched(monkeypatch): + calls = [] + + def fake_load_dataset(path, **kwargs): + calls.append((path, kwargs)) + + monkeypatch.setattr(prefetch, "load_dataset", fake_load_dataset) + + prefetch._prefetch_dataset( + DatasetEntry( + name="reasoning", + path="org/reasoning", + revision="deadbeef", + subset="math", + split="train", + ) + ) + + assert calls == [ + ( + "org/reasoning", + {"split": "train", "name": "math", "revision": "deadbeef"}, + ) + ] diff --git a/tests/test_slurm_render.py b/tests/test_slurm_render.py index a0da377..088f136 100644 --- a/tests/test_slurm_render.py +++ b/tests/test_slurm_render.py @@ -376,6 +376,38 @@ def test_trl_deepspeed_flags_absent_when_accelerate_override_disables(tmp_path, assert "--deepspeed_multinode_launcher" not in content +@pytest.mark.parametrize( + "renderer", + [render_trl_slurm_script, render_trl_container_slurm_script], +) +def test_trl_accelerate_config_is_frozen_and_forwarded(tmp_path, config, renderer): + profile = tmp_path / "fsdp.yaml" + profile.write_text("distributed_type: FSDP\n") + config.accelerate.config_file = str(profile) + config.deepspeed = None + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + content = renderer(config, run_dir, "configs/trl/sft.yaml").read_text() + frozen = run_dir / "slurm" / "accelerate.yaml" + + assert frozen.read_text() == "distributed_type: FSDP\n" + assert f"--config_file {frozen.resolve()}" in content + assert "--use_deepspeed" not in content + + profile.write_text("distributed_type: MULTI_GPU\n") + assert frozen.read_text() == "distributed_type: FSDP\n" + + +def test_missing_accelerate_config_fails_before_submission(tmp_path, config): + config.accelerate.config_file = str(tmp_path / "missing.yaml") + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + with pytest.raises(FileNotFoundError, match="accelerate.config_file"): + render_trl_slurm_script(config, run_dir, "configs/trl/sft.yaml") + + # --------------------------------------------------------------------------- # container.path — configurable PATH inside the container # ---------------------------------------------------------------------------