diff --git a/README.md b/README.md index f1405f8..11013a1 100644 --- a/README.md +++ b/README.md @@ -271,6 +271,8 @@ container: env_file: env/cluster.env # required when image is set; sourced before launch ``` +The env file must export `HF_HOME`, `HF_HUB_CACHE`, and `HUGGINGFACE_HUB_CACHE`. `scripts/submit.py` refuses to submit if any of them is unset, and the job checks again before launching the container. + Both the LlamaFactory and containerized TRL backends use this mechanism. Building containers for different HPCs is a work in progress, so if your cluster specific container is not available yet, please use the `uv` environment instead (or raise a pull request with a recipe for your cluster-specific container!). For TRL, use the `uv` environment by setting `container: null` or `container.image: null`; LlamaFactory requires a container image. ## 🧩 Feature Guide @@ -424,7 +426,7 @@ You must specify exactly one determining factor for training duration in the `tr - **Offline**: `offline: true` Disables Hugging Face Hub / Weights & Biases network calls (essential for air-gapped nodes). - **Prefetch assets**: `prefetch_assets: true` (default) - When submitting via `scripts/submit.py`, pre-fetches models and datasets on the login node before the job is submitted to warm the shared Hugging Face caches (useful for faster startup and for `offline: true` air-gapped runs). Runs whenever `offline` or `prefetch_assets` is true — set `prefetch_assets: false` to skip this step (e.g., assets are already warm in the shared cache). A tokenizer from a different repo or revision (`model.tokenizer_name_or_path` / `model.tokenizer_revision`) is fetched too, without its model weights. + When submitting via `scripts/submit.py`, pre-fetches models and datasets on the login node before the job is submitted to warm the shared Hugging Face caches (useful for faster startup and for `offline: true` air-gapped runs). Runs whenever `offline` or `prefetch_assets` is true — set `prefetch_assets: false` to skip this step (e.g., assets are already warm in the shared cache). A tokenizer from a different repo or revision (`model.tokenizer_name_or_path` / `model.tokenizer_revision`) is fetched too, without its model weights. For `backend: llamafactory`, it downloads `llamafactory.model_name_or_path` and `ref_model` (with `model_revision` and `cache_dir`), and loads the datasets named in `dataset` / `eval_dataset` as LlamaFactory does: `hf_hub_url` registry entries, `dataset_dir: ONLINE`, and `REMOTE:` registries. Sources it cannot download (`ms_hub_url`, `om_hub_url`, `cloud_file_name`, `script_url`, `streaming: true`) are rejected at submission when `offline: true`, and skipped otherwise. The local snapshot directories of the models are written into `llamafactory_config.yaml`; `config.yaml` keeps the Hub IDs. - **Debug**: `debug.enabled: true` Forces `report_to: none`, uses a separate output directory, and allows overwriting existing runs. - **Tokenize only**: `--tokenize-only` (CLI flag on `train.py` / `submit.py`) @@ -479,7 +481,7 @@ Each run generates a unique directory based on `paths.output_base` (or `paths.de ## LlamaFactory Backend -An alternative backend using [LlamaFactory](https://github.com/hiyouga/LlamaFactory) for training, running inside a Singularity container. +An alternative backend using [LlamaFactory](https://github.com/hiyouga/LlamaFactory) for training, running inside a Singularity container. The submission targets LlamaFactory **v0.9.5**: the stage choices and the input path keys come from that release, so the container must run it. ### Setup @@ -502,6 +504,22 @@ python scripts/submit.py --config configs/llamafactory/long-context.yaml - DeepSpeed: `configs/deepspeed/z3_partial_offload.json` - Dataset registry: `data/llamafactory/dataset_info.json` +### Configuration + +- Set every training setting in the `llamafactory` block, with LlamaFactory's own keys. The TRL sections (`model`, `training`, `checkpointing`, `data`, `sft`, `dpo`, `deepspeed`, `accelerate`, `logging`) stop the submission with an error, also when they are null or come from a CLI override. +- `llamafactory.stage` (default `sft`) must be one of `pt`, `sft`, `rm`, `ppo`, `dpo`, `kto`, and the top-level `method` must have the same value. For example, a DPO run sets `method: dpo` and `stage: dpo`. +- The submission review shows the values of the `llamafactory` block. Keys that the block does not set show the LlamaFactory or transformers default, marked `(default)`. +- The job always launches `llamafactory-cli train` through `torchrun` (`FORCE_TORCHRUN=1`), also with one GPU per node. +- LlamaFactory v0.9.5 does not resume the `rm` and `ppo` stages with `finetuning_type: full` or `freeze`. A requeued job of this type starts again from step 0, and the review shows a warning. + +### Run directory + +- The frozen `/config.yaml` holds only the `llamafactory` block and the submission keys, so you can submit again from it. +- The run name is `llamafactory----`, built from `llamafactory.stage` (default `sft`), `model_name_or_path`, and the names in `dataset` joined with `+`. Set `run_name` for a long dataset mix, because a directory name cannot exceed 255 bytes. +- `scripts/submit.py` copies the DeepSpeed JSON and `dataset_info.json` into `/llamafactory/`, and the job runs in ``. Later edits in the repository do not change a queued or resumed run. A new submit with the same `run_name` keeps the existing copies. The review screen and the log warn about each kept copy and say when it differs from the current file; to use the current file, delete the copy or set a new `run_name`. +- Relative paths in the `llamafactory` block start at the directory where you run `submit.py`. Existing input paths (models, adapters, `media_dir`, ...) are made absolute; output paths such as `export_dir` stay relative to ``. +- LlamaFactory writes checkpoints and the final model to `/checkpoints/`, and a requeued or resubmitted job resumes from the last checkpoint. Do not set `output_dir` or `overwrite_output_dir`; validation rejects them. + ## 📘 Configuration Reference: `configs/trl/sft.yaml` Full reference configuration for the default SFT setup: diff --git a/configs/llamafactory/long-context.yaml b/configs/llamafactory/long-context.yaml index 391229f..60a8542 100644 --- a/configs/llamafactory/long-context.yaml +++ b/configs/llamafactory/long-context.yaml @@ -59,10 +59,8 @@ llamafactory: group_by_length: true # output - output_dir: output/long-context logging_steps: 5 plot_loss: true - overwrite_output_dir: true save_only_model: false report_to: none diff --git a/configs/llamafactory/prelude-sft.yaml b/configs/llamafactory/prelude-sft.yaml new file mode 100644 index 0000000..f96dcd7 --- /dev/null +++ b/configs/llamafactory/prelude-sft.yaml @@ -0,0 +1,123 @@ +# ============================================================================ +# LlamaFactory SFT of Prelude 9B on OpenEuroLLM-Excellent-SFT (LUMI) +# ============================================================================ +# Usage: python scripts/submit.py --config configs/llamafactory/prelude-sft.yaml +# +# Needs LlamaFactory 0.9.5 in container.image. Before the first submission, make +# the model directory in llamafactory.model_name_or_path: the model weights of +# birgermoell/oellm-9b-256k-theta64m-prelude-anneal300b with the tokenizer of +# openeurollm/tokenizer-256k (revision qwen3-tokens). LlamaFactory loads the +# tokenizer from model_name_or_path, and only this tokenizer has the qwen3 tokens. +# ============================================================================ + +backend: llamafactory +method: sft +run_name: oellm-9b-256k-theta64m-prelude-anneal300b-excellent-sft-llamafactory +offline: false + +# -- Container --------------------------------------------------------------- +container: + image: /scratch/project_465002530/containers/REPLACE-WITH-llamafactory-0.9.5-rocm.sif # TODO: LlamaFactory 0.9.5 ROCm image + bind_mounts: + - /pfs/lustrep3/scratch/project_465002530/users/krishnak/post-training/ # replace with your own path to the post-training repo + path: /opt/venv/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin # check against the image + env_file: env/lumi.env # make sure you have this file in the post-training repo, or change to your own env file + +# -- SLURM ------------------------------------------------------------------- +slurm: # LUMI Cluster SLURM configuration + account: "project_465002530" + partition: "standard-g" # Test on dev-g first + num_nodes: 4 + gpus_per_node: 8 + cpus_per_task: 56 + wall_time: "36:00:00" + job_name: "Prelude-SFT" + signal_time_seconds: 300 # SIGUSR1 sent this many seconds before timeout to trigger self-healing + max_failures: 1 # Self-healing retry limit + mem: "256G" + +# -- Debug mode -------------------------------------------------------------- +debug: + enabled: false + override_existing: false + +# -- Output paths ------------------------------------------------------------- +paths: + output_base: "outputs" + debug_base: "outputs/debug" + +# -- LlamaFactory-native training config ------------------------------------ +llamafactory: + # model + model_name_or_path: /scratch/project_465002530/users/krishnak/models/oellm-9b-256k-theta64m-prelude-anneal300b-qwen3-tokens # TODO: see the header + flash_attn: fa2 # No flash attention 3 for rocm yet + enable_liger_kernel: true + disable_gradient_checkpointing: false + use_reentrant_gc: false + + # method + stage: sft + do_train: true + finetuning_type: full + deepspeed: + bf16: + enabled: auto + zero_optimization: + stage: 2 + overlap_comm: false # Enabling overlap_comm has caused issues before + contiguous_gradients: true + reduce_scatter: true + gradient_clipping: 1.0 + train_micro_batch_size_per_gpu: "auto" + gradient_accumulation_steps: "auto" + train_batch_size: "auto" + optimizer: + type: AdamW + params: + lr: "auto" + betas: "auto" + eps: "auto" + weight_decay: "auto" + + # dataset + dataset_dir: data/llamafactory + dataset: openeurollm_excellent_sft + template: qwen3 + enable_thinking: true # empty block on the trained answer + mask_history: true # loss on the final assistant message only + cutoff_len: 32768 + packing: true + neat_packing: true # no attention across packed examples (needs batch size 1) + preprocessing_num_workers: 32 + + # train + per_device_train_batch_size: 1 + gradient_accumulation_steps: 1 # 1 x 32 GPUs x 1 = 32 per step + num_train_epochs: 2.0 + learning_rate: 8.0e-5 + lr_scheduler_type: linear + warmup_steps: 0.03 # a value < 1 is a ratio (transformers >= 5); use warmup_ratio with transformers 4.x + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_epsilon: 1.0e-8 + weight_decay: 0.0 + max_grad_norm: 1.0 + bf16: true + seed: 42 + torch_compile: true + torch_compile_backend: inductor + + # checkpoints + save_strategy: steps + save_steps: 250 + save_total_limit: 3 + save_only_model: false + + # logging + run_name: oellm-9b-256k-theta64m-prelude-anneal300b-excellent-sft-llamafactory + report_to: + - wandb + - tensorboard + logging_dir: logs # relative to the run directory + logging_steps: 10 + include_num_input_tokens_seen: non_padding diff --git a/data/llamafactory/dataset_info.json b/data/llamafactory/dataset_info.json index 3a08599..4896016 100644 --- a/data/llamafactory/dataset_info.json +++ b/data/llamafactory/dataset_info.json @@ -12,5 +12,38 @@ "assistant_tag": "assistant", "system_tag": "system" } + }, + "dolci_instruct_sft": { + "hf_hub_url": "allenai/Dolci-Instruct-SFT", + "split": "train", + "formatting": "sharegpt", + "columns": { + "messages": "messages" + }, + "tags": { + "role_tag": "role", + "content_tag": "content", + "user_tag": "user", + "assistant_tag": "assistant", + "observation_tag": "environment", + "system_tag": "system" + } + }, + "openeurollm_excellent_sft": { + "hf_hub_url": "openeurollm/OpenEuroLLM-Excellent-SFT", + "split": "train", + "formatting": "sharegpt", + "columns": { + "messages": "messages", + "tools": "tools" + }, + "tags": { + "role_tag": "role", + "content_tag": "content", + "user_tag": "user", + "assistant_tag": "assistant", + "observation_tag": "tool", + "system_tag": "system" + } } } diff --git a/scripts/submit.py b/scripts/submit.py index 555241b..42fd374 100644 --- a/scripts/submit.py +++ b/scripts/submit.py @@ -23,7 +23,7 @@ sys.path.insert(0, str(_PROJECT_ROOT / "src")) from post_training.config import PostTrainingConfig -from post_training.slurm.launcher import generate_and_submit +from post_training.slurm.launcher import REQUIRED_HF_CACHE_VARS, generate_and_submit from post_training.utils.guardrails import run_guardrails from post_training.utils.logging import setup_logging from post_training.utils.paths import setup_run_directory @@ -75,6 +75,45 @@ def _apply_hf_env_from_file(env_file: str) -> None: logger.info("Applied HF cache vars from %s: %s", env_file, ", ".join(applied)) +def _check_required_hf_vars(config: PostTrainingConfig) -> None: + """Refuse a container job that would stop at start for a missing HF cache variable. + + Call after _apply_hf_env_from_file(): os.environ then holds the submitting + shell's values, which sbatch passes to the job, plus the env file's exports. + The job script checks again at start, for differences on the compute node. + """ + if config.container is None or not config.container.image: + return + missing = [var for var in REQUIRED_HF_CACHE_VARS if not os.environ.get(var)] + if missing: + raise ValueError( + f"{', '.join(missing)} not set. Export them in container.env_file " + f"({config.container.env_file}) as 'export NAME=value' lines." + ) + + +def _prepare_llamafactory_tokenize_only(config: PostTrainingConfig) -> None: + """Check llamafactory.tokenized_path for a --tokenize-only job and make it absolute. + + The job saves the tokenized data there. A relative path starts at the directory + where submit.py runs, so a later training submit with the same value finds the data. + """ + lf = config.llamafactory + tokenized_path = lf.get("tokenized_path") + if not isinstance(tokenized_path, str) or not tokenized_path: + raise ValueError( + "--tokenize-only with backend: llamafactory needs llamafactory.tokenized_path: " + "the job saves the tokenized data there." + ) + path = Path(tokenized_path).resolve() + if path.is_dir() and any(path.iterdir()): + raise ValueError( + f"llamafactory.tokenized_path '{path}' already holds data, and LlamaFactory would " + "load it. Delete it, or set a new tokenized_path." + ) + lf["tokenized_path"] = str(path) + + def _parse_args() -> tuple[str, list[str], bool]: parser = argparse.ArgumentParser(description="Submit a SLURM training job.") parser.add_argument( @@ -91,7 +130,10 @@ def _parse_args() -> tuple[str, list[str], bool]: parser.add_argument( "--tokenize-only", action="store_true", - help="Pass --tokenize-only to train.py — exits after trainer initialization.", + help=( + "Tokenize the datasets on 1 GPU and exit before training. TRL: passes " + "--tokenize-only to train.py. LlamaFactory: saves to llamafactory.tokenized_path." + ), ) known, unknown = parser.parse_known_args() return known.config, unknown, known.confirm, known.tokenize_only @@ -106,16 +148,15 @@ def main() -> None: if tokenize_only: if config.backend == "llamafactory": - raise ValueError( - "--tokenize-only is only supported with backend=trl " - "(LlamaFactory uses llamafactory-cli, which has no equivalent flag)." - ) + _prepare_llamafactory_tokenize_only(config) config.slurm.num_nodes = 1 config.slurm.gpus_per_node = 1 if config.container is not None and config.container.env_file: _apply_hf_env_from_file(config.container.env_file) + _check_required_hf_vars(config) + prefetched = None if config.offline or config.prefetch_assets: logger.info( @@ -144,7 +185,9 @@ def main() -> None: # the original repo id) are done. train.py then loads a local directory # path, so from_pretrained skips HF Hub's cache-resolution/filelock path # entirely — avoiding contention when every rank loads it at once. - if prefetched is not None: + # LlamaFactory gets its snapshot paths in post_freeze() instead, so + # config.yaml keeps the user's llamafactory block and TRL defaults. + if prefetched is not None and config.backend == "trl": config.model.name_or_path = prefetched.model config.model.tokenizer_name_or_path = prefetched.tokenizer if prefetched.ref_model is not None: @@ -157,7 +200,7 @@ def main() -> None: # Copy any backend-specific artifacts for reproducibility. from post_training.backend import get_backend - get_backend(config.backend).post_freeze(config, run_dir) + get_backend(config.backend).post_freeze(config, run_dir, prefetched) job_id = generate_and_submit( config, run_dir, str(frozen.resolve()), tokenize_only=tokenize_only diff --git a/src/post_training/backend.py b/src/post_training/backend.py index b3f4685..c4b2739 100644 --- a/src/post_training/backend.py +++ b/src/post_training/backend.py @@ -2,6 +2,7 @@ from __future__ import annotations +import copy import hashlib import json import logging @@ -15,7 +16,10 @@ import yaml if TYPE_CHECKING: + from collections.abc import Callable + from post_training.config import PostTrainingConfig + from post_training.utils.prefetch import PrefetchedPaths logger = logging.getLogger(__name__) @@ -50,8 +54,24 @@ def render_slurm_script( """Render the SLURM batch script and return its path.""" @abstractmethod - def post_freeze(self, config: PostTrainingConfig, run_dir: Path) -> None: - """Copy any extra artifacts into the run directory after config freeze.""" + def post_freeze( + self, + config: PostTrainingConfig, + run_dir: Path, + prefetched: PrefetchedPaths | None = None, + ) -> None: + """Copy any extra artifacts into the run directory after config freeze. + + *prefetched* holds the local snapshot directories from the login-node + prefetch, or None when the prefetch did not run. + """ + + def check_explicit_keys(self, explicit: dict) -> None: + """Check the keys that the YAML file and the CLI overrides set. + + *explicit* holds only those keys, before the schema defaults fill in. + The default accepts every key. + """ # --------------------------------------------------------------------------- @@ -60,6 +80,20 @@ def post_freeze(self, config: PostTrainingConfig, run_dir: Path) -> None: _SUPPORTED_METHODS = ("sft", "dpo") +# Config sections that only the TRL backend reads. A LlamaFactory run takes all +# training settings from its llamafactory block, so it rejects these. +TRL_SECTIONS = ( + "model", + "training", + "checkpointing", + "data", + "sft", + "dpo", + "deepspeed", + "accelerate", + "logging", +) + def _shorten_model_name(name_or_path: str) -> str: short = name_or_path.rsplit("/", 1)[-1] @@ -190,9 +224,12 @@ def render_slurm_script(self, config, run_dir, frozen_config_path, *, tokenize_o config, run_dir, frozen_config_path, tokenize_only=tokenize_only ) - def post_freeze(self, config, run_dir): + def post_freeze(self, config, run_dir, prefetched=None): """Freeze src/post_training and scripts/ into run_dir for containerized runs. + *prefetched* is not used: submit.py writes the TRL snapshot paths into + config.yaml itself. + Done here (at submission time) rather than inside the SLURM job script, so a job that sits queued isn't affected by later repo changes. Each directory is copied into a temp path and atomically swapped into place, @@ -224,6 +261,141 @@ def _freeze_dir(src: Path, dst: Path) -> None: # LlamaFactory # --------------------------------------------------------------------------- +# The LlamaFactory release that the container runs. The key names and stage choices +# below come from its hparams. +LLAMAFACTORY_VERSION = "0.9.5" + +# llamafactory keys that hold local input paths. The job runs in the run directory, so a +# relative value that exists is made absolute at submission. Output keys (export_dir, +# logging_dir, offload_folder, swanlab_logdir, profiler_output_dir) stay relative and land +# in the run directory. +_LLAMAFACTORY_INPUT_PATH_KEYS = ( + "model_name_or_path", + "adapter_folder", + "cache_dir", + "new_special_tokens_config", + "export_quantization_dataset", + "kt_weight_path", + "kt_expert_checkpoint_path", + "hyper_parallel_args", + "media_dir", + "tokenized_path", + "ref_model", + "reward_model", + "resume_from_checkpoint", + "fsdp_config", +) +# Keys that hold a comma-separated list of input paths. +_LLAMAFACTORY_INPUT_PATH_LIST_KEYS = ( + "adapter_name_or_path", + "ref_model_adapters", + "reward_model_adapters", +) + +_LLAMAFACTORY_DEFAULT_DATASET_DIR = "data" +_LLAMAFACTORY_DEFAULT_STAGE = "sft" +_LLAMAFACTORY_STAGES = ("pt", "sft", "rm", "ppo", "dpo", "kto") +# LlamaFactory does not resume these stages with full or freeze tuning: a requeued job +# starts again from step 0. +_LLAMAFACTORY_NON_RESUMABLE_STAGES = ("rm", "ppo") +_LLAMAFACTORY_REGISTRY = "dataset_info.json" +# A --tokenize-only job runs this script in the container, not llamafactory-cli train. +_LLAMAFACTORY_TOKENIZE_SCRIPT = ( + Path(__file__).resolve().parent / "slurm" / "llamafactory_tokenize.py" +) + + +def llamafactory_dataset_names( + lf: dict, keys: tuple[str, ...] = ("dataset", "eval_dataset") +) -> list[str]: + """The dataset names under *keys*, split like LlamaFactory's split_arg().""" + names: list[str] = [] + for key in keys: + value = lf.get(key) + if isinstance(value, str): + value = value.split(",") + if isinstance(value, list): + names.extend(str(name).strip() for name in value if str(name).strip()) + return names + + +def _is_local_registry(dataset_dir: str) -> bool: + """False for the dataset_dir values that LlamaFactory resolves on the Hub.""" + return dataset_dir != "ONLINE" and not dataset_dir.startswith("REMOTE:") + + +def _absolute_if_exists(value: str) -> str: + """Return a relative path that exists as an absolute path; return anything else unchanged.""" + path = Path(value) + if path.is_absolute() or not path.exists(): + return value + return str(path.resolve()) + + +def _frozen_copy_differs(dst: Path, source: Path, render: Callable[[], str]) -> bool: + """True when the frozen *dst* is not what the current *source* would freeze to.""" + return source.is_file() and dst.read_text() != render() + + +def _freeze_file(dst: Path, source: Path, render: Callable[[], str]) -> None: + """Write render() to *dst* once; an existing *dst* is kept, so a resumed run keeps its inputs. + + The text goes to a temp file first and is swapped into place, so an interrupted + submission never leaves a partial copy that a later one would keep. + """ + if dst.exists(): + if _frozen_copy_differs(dst, source, render): + logger.warning( + "%s differs from %s — keeping the frozen copy. Delete it, or use a new " + "run_name, to freeze the current file.", + dst, + source, + ) + else: + logger.warning("%s already exists — keeping the frozen copy.", dst) + return + dst.parent.mkdir(parents=True, exist_ok=True) + tmp = dst.with_name(dst.name + ".tmp") + tmp.write_text(render()) + os.replace(tmp, dst) + + +def _registry_pointing_at(source_dir: Path) -> str: + """Return source_dir's dataset_info.json with relative local entries made absolute.""" + entries = json.loads((source_dir / _LLAMAFACTORY_REGISTRY).read_text()) + for entry in entries.values(): + for key in ("file_name", "script_url"): + if isinstance(entry.get(key), str) and not Path(entry[key]).is_absolute(): + entry[key] = str(source_dir / entry[key]) + return json.dumps(entries, indent=2, ensure_ascii=False) + "\n" + + +def _llamafactory_frozen_inputs( + lf: dict, frozen_dir: Path +) -> dict[str, tuple[Path, Path, Callable[[], str]]]: + """The files that post_freeze() freezes, as {kind: (frozen copy, source, render)}. + + Relative sources start at the directory where submit.py runs. + """ + inputs: dict[str, tuple[Path, Path, Callable[[], str]]] = {} + # Frozen for every submission, so a later --tokenize-only submit into the run + # directory runs the same script; only a --tokenize-only job uses it. + source = _LLAMAFACTORY_TOKENIZE_SCRIPT + inputs["tokenize"] = (frozen_dir / source.name, source, source.read_text) + deepspeed = lf.get("deepspeed") + if isinstance(deepspeed, str): + source = Path(deepspeed) + inputs["deepspeed"] = (frozen_dir / source.name, source, source.read_text) + dataset_dir = lf.get("dataset_dir") or _LLAMAFACTORY_DEFAULT_DATASET_DIR + if _is_local_registry(dataset_dir): + source_dir = Path(dataset_dir).resolve() + inputs["registry"] = ( + frozen_dir / _LLAMAFACTORY_REGISTRY, + source_dir / _LLAMAFACTORY_REGISTRY, + lambda: _registry_pointing_at(source_dir), + ) + return inputs + class LlamaFactoryBackend(Backend): def validate(self, config: PostTrainingConfig) -> None: @@ -231,24 +403,159 @@ def validate(self, config: PostTrainingConfig) -> None: raise ValueError("llamafactory must be set when backend='llamafactory'.") if config.container is None or not config.container.image: raise ValueError("container.image must be set when backend='llamafactory'.") + if not config.llamafactory.get("model_name_or_path"): + raise ValueError("llamafactory.model_name_or_path must be set.") + + lf = config.llamafactory + stage = lf.get("stage", _LLAMAFACTORY_DEFAULT_STAGE) + if stage not in _LLAMAFACTORY_STAGES: + raise ValueError( + f"llamafactory.stage '{stage}' is not a LlamaFactory {LLAMAFACTORY_VERSION} " + f"stage. Choose from: {', '.join(_LLAMAFACTORY_STAGES)}." + ) + if config.method != stage: + raise ValueError( + f"method '{config.method}' does not match llamafactory.stage '{stage}'. " + f"Set method: {stage}." + ) + + # post_freeze() owns these two settings, so a requeued or resubmitted job resumes. + if "output_dir" in lf: + raise ValueError( + "llamafactory.output_dir must not be set: the submission writes checkpoints to " + "/checkpoints so that a requeued job resumes. Remove it." + ) + if lf.get("overwrite_output_dir"): + raise ValueError( + "llamafactory.overwrite_output_dir: true stops a requeued job from resuming " + "from its checkpoints. Remove it." + ) + + # post_freeze() copies these files, so a missing one must fail before submission. + deepspeed = lf.get("deepspeed") + if isinstance(deepspeed, str) and not Path(deepspeed).is_file(): + raise ValueError( + f"llamafactory.deepspeed file '{deepspeed}' not found. Relative paths start at " + "the directory where submit.py runs." + ) + dataset_dir = lf.get("dataset_dir") or _LLAMAFACTORY_DEFAULT_DATASET_DIR + if (lf.get("dataset") or lf.get("eval_dataset")) and _is_local_registry(dataset_dir): + registry = Path(dataset_dir) / _LLAMAFACTORY_REGISTRY + if not registry.is_file(): + raise ValueError( + f"{registry} not found: named datasets need a {_LLAMAFACTORY_REGISTRY} in " + "llamafactory.dataset_dir. Relative paths start at the directory where " + "submit.py runs." + ) + + def check_explicit_keys(self, explicit: dict) -> None: + """Reject TRL sections, including null values and CLI overrides.""" + trl_keys: list[str] = [] + for section in TRL_SECTIONS: + if section not in explicit: + continue + value = explicit[section] + if isinstance(value, dict) and value: + trl_keys.extend(f"{section}.{key}" for key in value) + else: + trl_keys.append(section) + if trl_keys: + raise ValueError( + "backend: llamafactory takes all training settings from the llamafactory " + f"block. Remove these TRL keys: {', '.join(trl_keys)}." + ) def generate_run_name(self, config: PostTrainingConfig, timestamp: str) -> str: - return f"llamafactory-{config.method}-{timestamp}" + """llamafactory----. + + Built from the llamafactory block only; the TRL key `method` is not used. + """ + lf = config.llamafactory + parts = [ + "llamafactory", + str(lf.get("stage") or _LLAMAFACTORY_DEFAULT_STAGE), + _shorten_model_name(str(lf["model_name_or_path"]).rstrip("/")), + ] + names = llamafactory_dataset_names(lf, keys=("dataset",)) + if names: + parts.append("+".join(_shorten_dataset_name(name) for name in names)) + parts.append(timestamp) + return "-".join(parts) def run_dir_subdirs(self) -> list[str]: - return [] + return ["checkpoints"] + + def kept_frozen_inputs(self, config, run_dir) -> list[tuple[Path, Path, bool]]: + """(frozen copy, source, differs) for each frozen input that post_freeze() keeps. + + A new submit into an existing run directory keeps these copies and does not + read the current files, so the review screen lists them. + """ + frozen_dir = run_dir.resolve() / "llamafactory" + inputs = _llamafactory_frozen_inputs(config.llamafactory, frozen_dir) + return [ + (frozen, source, _frozen_copy_differs(frozen, source, render)) + for frozen, source, render in inputs.values() + if frozen.exists() + ] def render_slurm_script(self, config, run_dir, frozen_config_path, *, tokenize_only=False): - # LlamaFactory has no --tokenize-only equivalent; submit.py rejects it before this point. - del tokenize_only from post_training.slurm.launcher import render_llamafactory_slurm_script - return render_llamafactory_slurm_script(config, run_dir) + return render_llamafactory_slurm_script(config, run_dir, tokenize_only=tokenize_only) + + def post_freeze(self, config, run_dir, prefetched=None): + """Freeze the job's inputs into run_dir and write llamafactory_config.yaml. + + The DeepSpeed JSON, dataset_info.json and the tokenize script are copied into + run_dir/llamafactory/, so a queued or resumed job does not read later edits. Only + llamafactory_config.yaml carries the rewritten values; config.yaml keeps the + user's block, so a new submit from it passes validation again. + """ + run_dir = run_dir.resolve() + frozen_dir = run_dir / "llamafactory" + lf_config = copy.deepcopy(config.llamafactory) + + # Every rank then loads a plain local directory, not a Hub ID through the cache locks. + if prefetched is not None: + lf_config["model_name_or_path"] = prefetched.model + if prefetched.ref_model is not None: + lf_config["ref_model"] = prefetched.ref_model + + for key in _LLAMAFACTORY_INPUT_PATH_KEYS: + if isinstance(lf_config.get(key), str): + lf_config[key] = _absolute_if_exists(lf_config[key]) + for key in _LLAMAFACTORY_INPUT_PATH_LIST_KEYS: + if isinstance(lf_config.get(key), str): + parts = lf_config[key].split(",") + lf_config[key] = ",".join(_absolute_if_exists(part.strip()) for part in parts) + + inputs = _llamafactory_frozen_inputs(lf_config, frozen_dir) + if "tokenize" in inputs: + _freeze_file(*inputs["tokenize"]) + if "deepspeed" in inputs: + frozen_deepspeed, source, render = inputs["deepspeed"] + _freeze_file(frozen_deepspeed, source, render) + lf_config["deepspeed"] = str(frozen_deepspeed) + + if "registry" in inputs: + frozen_registry, source, render = inputs["registry"] + if source.is_file() or frozen_registry.exists(): + _freeze_file(frozen_registry, source, render) + lf_config["dataset_dir"] = str(frozen_dir) + # media_dir defaults to dataset_dir; keep relative media paths on the source. + lf_config.setdefault("media_dir", str(source.parent)) + elif source.parent.exists(): + lf_config["dataset_dir"] = str(source.parent) + + # LlamaFactory resumes from the last checkpoint-N in output_dir only when + # overwrite_output_dir is false. + lf_config["output_dir"] = str(run_dir / "checkpoints") + lf_config["overwrite_output_dir"] = False - def post_freeze(self, config, run_dir): lf_config_path = run_dir / "llamafactory_config.yaml" with open(lf_config_path, "w") as f: - yaml.dump(config.llamafactory, f, default_flow_style=False) + yaml.dump(lf_config, f, default_flow_style=False) # --------------------------------------------------------------------------- diff --git a/src/post_training/config.py b/src/post_training/config.py index 6a20f1a..1d4b7c5 100644 --- a/src/post_training/config.py +++ b/src/post_training/config.py @@ -341,12 +341,20 @@ def load( file_cfg = OmegaConf.create(file_dict) merged: DictConfig = OmegaConf.merge(schema, file_cfg) + explicit_cfg = file_cfg if cli_overrides: cli_cfg = OmegaConf.from_dotlist(cli_overrides) merged = OmegaConf.merge(merged, cli_cfg) + explicit_cfg = OmegaConf.merge(file_cfg, cli_cfg) config: PostTrainingConfig = OmegaConf.to_object(merged) # type: ignore[assignment] + + from post_training.backend import get_backend + + # The keys the user wrote, in the file or on the CLI, before defaults fill in. + explicit = OmegaConf.to_container(explicit_cfg, resolve=False) + get_backend(config.backend).check_explicit_keys(explicit) config._validate() return config @@ -354,8 +362,21 @@ def load( # Saving # ------------------------------------------------------------------ def save(self, yaml_path: str | Path) -> None: - """Save the configuration to a YAML file.""" - OmegaConf.save(self, yaml_path) + """Save the configuration to a YAML file. + + A LlamaFactory config is saved without the TRL sections: they have no effect + on it, and load() rejects them. + """ + if self.backend != "llamafactory": + OmegaConf.save(self, yaml_path) + return + + from post_training.backend import TRL_SECTIONS + + frozen = OmegaConf.to_container(OmegaConf.structured(self)) + for section in TRL_SECTIONS: + frozen.pop(section, None) + OmegaConf.save(OmegaConf.create(frozen), yaml_path) # ------------------------------------------------------------------ # Validation & derived values diff --git a/src/post_training/slurm/job_llamafactory.sh.jinja b/src/post_training/slurm/job_llamafactory.sh.jinja index edd99cc..558d470 100644 --- a/src/post_training/slurm/job_llamafactory.sh.jinja +++ b/src/post_training/slurm/job_llamafactory.sh.jinja @@ -30,9 +30,31 @@ set -euo pipefail +{% if modules %} +# ── Environment modules ──────────────────────────────────────────────────── +{% if module_purge %} +module purge +{% endif %} +{% for mod in modules %} +module load {{ mod }} +{% endfor %} +{% endif %} + # ── Source cluster environment ──────────────────────────────────────────── {% if env_file %}source "{{ env_file }}"{% endif %} +# ── Check the HF cache variables ────────────────────────────────────────── +MISSING_HF_VARS="" +for var in {{ required_hf_vars | join(" ") }}; do + if [ -z "${!var:-}" ]; then + MISSING_HF_VARS="$MISSING_HF_VARS $var" + fi +done +if [ -n "$MISSING_HF_VARS" ]; then + echo "ERROR:$MISSING_HF_VARS not set. Export them in container.env_file." >&2 + exit 1 +fi + # ── Distributed setup ───────────────────────────────────────────────────── MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) MASTER_ADDR=$(getent hosts "$MASTER_ADDR" | awk '{print $1}') @@ -44,7 +66,8 @@ export WORLD_SIZE=$(( SLURM_NNODES * GPUS_PER_NODE )) # NCCL tuning for multi-node stability export NCCL_IB_TIMEOUT=120 export NCCL_DEBUG=WARN -export CUDA_DEVICE_MAX_CONNECTIONS=1 +# CUDA_DEVICE_MAX_CONNECTIONS=1 is a Megatron-LM tensor-parallel flag that +# serializes CUDA streams and hurts ZeRO-2 overlap_comm — do not set it. export OMP_NUM_THREADS=1 echo "==========================================" @@ -63,9 +86,14 @@ if [ ! -f "$FAILURE_COUNT_FILE" ]; then echo 0 > "$FAILURE_COUNT_FILE" fi +REQUEUED=0 handle_signal() { echo "[$(date)] Received SIGUSR1 — timeout approaching. Requeuing..." - scontrol requeue $SLURM_JOB_ID + if ! scontrol requeue "$SLURM_JOB_ID"; then + echo "[$(date)] scontrol requeue failed. Exiting." >&2 + exit 1 + fi + REQUEUED=1 } trap handle_signal USR1 @@ -79,30 +107,55 @@ srun --export=ALL --wait=60 --kill-on-bad-exit=1 \ # Prevent host Python/PATH from interfering with container export PATH=\"{{ container_path }}\" - export PYTHONPATH=\"{{ repo_dir }}/src\" + unset PYTHONPATH export PYTHONNOUSERSITE=1 export NODE_RANK=\"\$SLURM_NODEID\" export NNODES=\"\$SLURM_NNODES\" export NPROC_PER_NODE=${GPUS_PER_NODE} + # llamafactory-cli starts torchrun only when it sees more than one local GPU, so a + # job with one GPU per node would otherwise train each node on its own. + export FORCE_TORCHRUN=1 - # Forward HF cache/offline settings from host env + # Forward HF cache paths from host env export HF_HOME=${HF_HOME} export HF_HUB_CACHE=${HF_HUB_CACHE} export HUGGINGFACE_HUB_CACHE=${HUGGINGFACE_HUB_CACHE} export HF_DATASETS_CACHE=${HF_DATASETS_CACHE:-${HF_HOME}/datasets} - export HF_DATASETS_OFFLINE=1 - export TRANSFORMERS_OFFLINE=1 - export HF_HUB_OFFLINE=1 export SSL_CERT_FILE=/etc/ssl/certs/ca-certificates.crt - cd {{ repo_dir }} - + # LlamaFactory jobs do not run scripts/train.py, so set offline mode here, after the env file. +{% if offline %} + export HF_HUB_OFFLINE=1 + export HF_DATASETS_OFFLINE=1 + export TRANSFORMERS_OFFLINE=1 + export WANDB_MODE=offline +{% else %} + export HF_HUB_OFFLINE=0 + export HF_DATASETS_OFFLINE=0 + export TRANSFORMERS_OFFLINE=0 +{% endif %} + + cd \"{{ run_dir }}\" + +{% if tokenize_only %} + # --tokenize-only: save the tokenized datasets to tokenized_path and exit before the model load. + # PYTHONSAFEPATH keeps the script directory off sys.path, so no file there shadows a module. + PYTHONSAFEPATH=1 torchrun --standalone --nproc_per_node 1 {{ tokenize_script }} {{ llamafactory_config }} +{% else %} llamafactory-cli train {{ llamafactory_config }} +{% endif %} " & -wait $! -EXIT_CODE=$? +# Capture the srun exit code without letting set -e exit before the failure handler. +EXIT_CODE=0 +wait $! || EXIT_CODE=$? + +# A requeued job restarts on its own; a resubmit as well would start a second job. +if [ "$REQUEUED" -eq 1 ]; then + echo "[$(date)] Job requeued. Skipping the failure handler." + exit 0 +fi # Handle failure / resubmit if [ $EXIT_CODE -ne 0 ]; then @@ -112,6 +165,8 @@ if [ $EXIT_CODE -ne 0 ]; then if [ $FAILURES -lt $MAX_FAILURES ]; then echo "[$(date)] Failure $FAILURES/$MAX_FAILURES. Resubmitting..." sbatch "$0" + # Report this failure to SLURM even though a new job was submitted. + exit $EXIT_CODE else echo "[$(date)] Max failures reached ($MAX_FAILURES). Stopping." exit 1 diff --git a/src/post_training/slurm/job_trl_container.sh.jinja b/src/post_training/slurm/job_trl_container.sh.jinja index 989a947..6b33741 100644 --- a/src/post_training/slurm/job_trl_container.sh.jinja +++ b/src/post_training/slurm/job_trl_container.sh.jinja @@ -43,6 +43,18 @@ module load {{ mod }} # ── Source cluster environment ──────────────────────────────────────────── {% if env_file %}source "{{ env_file }}"{% endif %} +# ── Check the HF cache variables ────────────────────────────────────────── +MISSING_HF_VARS="" +for var in {{ required_hf_vars | join(" ") }}; do + if [ -z "${!var:-}" ]; then + MISSING_HF_VARS="$MISSING_HF_VARS $var" + fi +done +if [ -n "$MISSING_HF_VARS" ]; then + echo "ERROR:$MISSING_HF_VARS not set. Export them in container.env_file." >&2 + exit 1 +fi + # ── Distributed setup ───────────────────────────────────────────────────── MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) MASTER_ADDR=$(getent hosts "$MASTER_ADDR" | awk '{print $1}') @@ -75,9 +87,14 @@ if [ ! -f "$FAILURE_COUNT_FILE" ]; then echo 0 > "$FAILURE_COUNT_FILE" fi +REQUEUED=0 handle_signal() { echo "[$(date)] Received SIGUSR1 — timeout approaching. Requeuing..." - scontrol requeue $SLURM_JOB_ID + if ! scontrol requeue "$SLURM_JOB_ID"; then + echo "[$(date)] scontrol requeue failed. Exiting." >&2 + exit 1 + fi + REQUEUED=1 } trap handle_signal USR1 @@ -128,8 +145,15 @@ srun --export=ALL --wait=60 --kill-on-bad-exit=1 \ scripts/train.py --config {{ config_path }}{% if tokenize_only %} --tokenize-only{% endif %} " & -wait $! -EXIT_CODE=$? +# Capture the srun exit code without letting set -e exit before the failure handler. +EXIT_CODE=0 +wait $! || EXIT_CODE=$? + +# A requeued job restarts on its own; a resubmit as well would start a second job. +if [ "$REQUEUED" -eq 1 ]; then + echo "[$(date)] Job requeued. Skipping the failure handler." + exit 0 +fi # ── Handle failure / resubmit ───────────────────────────────────────────── if [ $EXIT_CODE -ne 0 ]; then @@ -139,6 +163,8 @@ if [ $EXIT_CODE -ne 0 ]; then if [ $FAILURES -lt $MAX_FAILURES ]; then echo "[$(date)] Failure $FAILURES/$MAX_FAILURES. Resubmitting..." sbatch "$0" + # Report this failure to SLURM even though a new job was submitted. + exit $EXIT_CODE else echo "[$(date)] Max failures reached ($MAX_FAILURES). Stopping." exit 1 diff --git a/src/post_training/slurm/launcher.py b/src/post_training/slurm/launcher.py index fd8338f..950b5cb 100644 --- a/src/post_training/slurm/launcher.py +++ b/src/post_training/slurm/launcher.py @@ -24,6 +24,10 @@ _LLAMAFACTORY_TEMPLATE_NAME = "job_llamafactory.sh.jinja" _TRL_CONTAINER_TEMPLATE_NAME = "job_trl_container.sh.jinja" +# HF cache variables that the container job scripts forward into the container. +# scripts/submit.py checks them at submission, and the job script checks them again at start. +REQUIRED_HF_CACHE_VARS = ("HF_HOME", "HF_HUB_CACHE", "HUGGINGFACE_HUB_CACHE") + def render_trl_slurm_script( config: PostTrainingConfig, @@ -144,6 +148,7 @@ def render_trl_container_slurm_script( container_image=config.container.image, bind_mounts=config.container.bind_mounts, env_file=config.container.env_file, + required_hf_vars=REQUIRED_HF_CACHE_VARS, container_path=config.container.path, ) @@ -160,8 +165,14 @@ def render_trl_container_slurm_script( def render_llamafactory_slurm_script( config: PostTrainingConfig, run_dir: Path, + *, + tokenize_only: bool = False, ) -> Path: - """Render the LlamaFactory SLURM batch script into *run_dir/slurm/job.sh*.""" + """Render the LlamaFactory SLURM batch script into *run_dir/slurm/job.sh*. + + With *tokenize_only*, the job runs run_dir/llamafactory/llamafactory_tokenize.py, which + post_freeze() froze, and not ``llamafactory-cli train``. + """ if config.container is None or not config.container.image: raise ValueError("container.image must be set to render a LlamaFactory script.") @@ -185,15 +196,20 @@ def render_llamafactory_slurm_script( wall_time=config.slurm.wall_time, signal_time_seconds=config.slurm.signal_time_seconds, max_failures=config.slurm.max_failures, + modules=config.slurm.modules, + module_purge=config.slurm.module_purge, + offline=config.offline, run_dir=str(run_dir.resolve()), # Container container_image=config.container.image, bind_mounts=config.container.bind_mounts, env_file=config.container.env_file, + required_hf_vars=REQUIRED_HF_CACHE_VARS, container_path=config.container.path, # LlamaFactory - llamafactory_config=str(run_dir / "llamafactory_config.yaml"), - repo_dir=str(Path.cwd()), + llamafactory_config=str(run_dir.resolve() / "llamafactory_config.yaml"), + tokenize_only=tokenize_only, + tokenize_script=str(run_dir.resolve() / "llamafactory" / "llamafactory_tokenize.py"), ) slurm_dir = run_dir / "slurm" diff --git a/src/post_training/slurm/llamafactory_tokenize.py b/src/post_training/slurm/llamafactory_tokenize.py new file mode 100644 index 0000000..50033a4 --- /dev/null +++ b/src/post_training/slurm/llamafactory_tokenize.py @@ -0,0 +1,108 @@ +"""Tokenize the datasets of a LlamaFactory config into its tokenized_path, without training. + +The submission copies this file into the run directory, and a --tokenize-only job runs it +under torchrun inside the LlamaFactory container. post_training does not import it: the +tokenize step needs the container's llamafactory package, so that import is deferred to +tokenize(). + +It runs the first steps of the LlamaFactory 0.9.5 workflows (load_tokenizer, +get_template_and_fix_tokenizer, get_dataset) and exits before load_model. + +Two jobs can target the same tokenized_path, e.g. after a double submit, so each job +writes to its own partial directory and publishes it with one rename. The rename of a +directory fails when the target is a non-empty directory, so the first job to publish +wins and the other discards its data; neither can expose a half-written directory. +""" + +import errno +import os +import shutil +import sys + +import yaml + +# The stage that each workflow passes to get_dataset(); DPO reads pairwise data like RM. +_DATASET_STAGE = {"pt": "pt", "sft": "sft", "rm": "rm", "dpo": "rm", "kto": "kto", "ppo": "ppo"} + + +def has_data(path: str) -> bool: + """Same test as LlamaFactory's has_tokenized_data(): a directory that is not empty.""" + return os.path.isdir(path) and bool(os.listdir(path)) + + +def partial_path(tokenized_path: str) -> str: + """The directory that only this job writes to, named after its SLURM job ID. + + A requeued job keeps its job ID, so it finds and removes its own interrupted save. + Outside SLURM the process ID keeps two local runs apart. + """ + job_id = os.environ.get("SLURM_JOB_ID") or str(os.getpid()) + return f"{tokenized_path.rstrip('/')}.partial-{job_id}" + + +def publish(partial: str, tokenized_path: str) -> bool: + """Rename *partial* to *tokenized_path*; False when another job published first. + + rename() replaces a missing or empty target directory in one step and fails on a + non-empty one, so two jobs cannot both publish. The loser removes its partial + directory, so nothing is left next to the data that the training job loads. + """ + try: + os.rename(partial, tokenized_path) + except OSError as exc: + if exc.errno not in (errno.ENOTEMPTY, errno.EEXIST): + raise + shutil.rmtree(partial) + return False + return True + + +def tokenize(args: dict) -> None: + """Run the LlamaFactory data pipeline for *args* and save it to args["tokenized_path"].""" + from llamafactory.data import get_dataset, get_template_and_fix_tokenizer + from llamafactory.hparams import get_train_args + from llamafactory.model import load_tokenizer + + model_args, data_args, training_args, finetuning_args, _ = get_train_args(args) + tokenizer_module = load_tokenizer(model_args) + template = get_template_and_fix_tokenizer(tokenizer_module["tokenizer"], data_args) + get_dataset( + template, + model_args, + data_args, + training_args, + stage=_DATASET_STAGE[finetuning_args.stage], + **tokenizer_module, + ) + + +def main(config_path: str) -> None: + with open(config_path) as f: + args = yaml.safe_load(f) + + # DeepSpeed does not change the tokenized data, and the job has one GPU. + args.pop("deepspeed", None) + + tokenized_path = args["tokenized_path"] + # Exit code 0: the job script would otherwise resubmit a job that has nothing to do. + if has_data(tokenized_path): + print(f"tokenized_path '{tokenized_path}' already holds data. Nothing to tokenize.") + return + + partial = partial_path(tokenized_path) + # Remove any interrupted save from a previous attempt so this run starts cleanly. + shutil.rmtree(partial, ignore_errors=True) + args["tokenized_path"] = partial + tokenize(args) + + if publish(partial, tokenized_path): + print(f"Tokenized dataset saved at {tokenized_path}. Train with this tokenized_path.") + else: + print( + f"Another job saved tokenized data at {tokenized_path} while this job ran. " + "The data of this job was discarded; train with that tokenized_path." + ) + + +if __name__ == "__main__": + main(sys.argv[1]) diff --git a/src/post_training/utils/guardrails.py b/src/post_training/utils/guardrails.py index 32ab111..0e90691 100644 --- a/src/post_training/utils/guardrails.py +++ b/src/post_training/utils/guardrails.py @@ -172,6 +172,16 @@ def run_guardrails(config: PostTrainingConfig, run_dir: Path, tokenize_only: boo _row("Container", "none (bare-metal)") _row("Offline mode", str(config.offline)) + if config.backend == "llamafactory": + _llamafactory_sections(config, run_dir, total_gpus, tokenize_only) + else: + _trl_sections(config, run_dir, total_gpus, tokenize_only) + _confirm(total_gpus) + + +def _trl_sections( + config: PostTrainingConfig, run_dir: Path, total_gpus: int, tokenize_only: bool +) -> None: # ------------------------------------------------------------------ # Model & data # ------------------------------------------------------------------ @@ -308,6 +318,132 @@ def run_guardrails(config: PostTrainingConfig, run_dir: Path, tokenize_only: boo warn=True, ) + +def _llamafactory_sections( + config: PostTrainingConfig, run_dir: Path, total_gpus: int, tokenize_only: bool +) -> None: + """Review the llamafactory block only; the TRL sections do not apply to these runs.""" + from post_training.backend import ( + _LLAMAFACTORY_DEFAULT_DATASET_DIR, + _LLAMAFACTORY_DEFAULT_STAGE, + _LLAMAFACTORY_NON_RESUMABLE_STAGES, + LLAMAFACTORY_VERSION, + LlamaFactoryBackend, + llamafactory_dataset_names, + ) + + lf = config.llamafactory or {} + stage = lf.get("stage", _LLAMAFACTORY_DEFAULT_STAGE) + finetuning_type = lf.get("finetuning_type", _LF_DEFAULTS["finetuning_type"]) + + _section("Model & Data") + _row("Model", str(lf.get("model_name_or_path"))) + if lf.get("ref_model"): + _row("Ref model", str(lf["ref_model"])) + _row("Template", _lf_value(lf, "template"), warn=lf.get("template") is None) + _row("Dataset dir", str(lf.get("dataset_dir") or _LLAMAFACTORY_DEFAULT_DATASET_DIR)) + _row("Datasets", ", ".join(llamafactory_dataset_names(lf, keys=("dataset",))) or "none") + eval_names = llamafactory_dataset_names(lf, keys=("eval_dataset",)) + if eval_names: + _row("Eval datasets", ", ".join(eval_names)) + _row("Cutoff length", _lf_value(lf, "cutoff_len")) + + _section(f"Training ({str(stage).upper()})") + _row("Backend", f"llamafactory {LLAMAFACTORY_VERSION}") + _row("Finetuning type", _lf_value(lf, "finetuning_type")) + if stage in ("dpo", "kto"): + _row("Preference loss", _lf_value(lf, "pref_loss")) + _row("Beta", _lf_value(lf, "pref_beta")) + max_steps = lf.get("max_steps", -1) + if isinstance(max_steps, int) and max_steps > 0: + _row("Duration", f"{max_steps:,} steps") + else: + _row("Epochs", _lf_value(lf, "num_train_epochs")) + _row("Learning rate", _lf_value(lf, "learning_rate")) + _row("LR scheduler", _lf_value(lf, "lr_scheduler_type")) + per_device = lf.get("per_device_train_batch_size", _LF_DEFAULTS["per_device_train_batch_size"]) + grad_acc = lf.get("gradient_accumulation_steps", _LF_DEFAULTS["gradient_accumulation_steps"]) + _row( + "Batch sizes", + f"{per_device} per device × {total_gpus} GPUs × {grad_acc} grad-acc" + f" = {per_device * total_gpus * grad_acc} effective", + ) + deepspeed = lf.get("deepspeed") + _row( + "DeepSpeed", "inline dict" if isinstance(deepspeed, dict) else str(deepspeed or "disabled") + ) + + _section("Checkpointing") + save_strategy = lf.get("save_strategy", _LF_DEFAULTS["save_strategy"]) + if save_strategy == "steps": + _row("Save steps", _lf_value(lf, "save_steps")) + else: + _row("Save strategy", str(save_strategy), warn=save_strategy == "no") + if stage in _LLAMAFACTORY_NON_RESUMABLE_STAGES and finetuning_type in ("full", "freeze"): + _row( + "Resume", + _yellow( + f"*** LlamaFactory cannot resume stage {stage} with {finetuning_type} tuning — " + "a requeued job starts again from step 0 ***" + ), + warn=True, + ) + elif save_strategy == "no": + _row("Resume", _yellow("*** no checkpoints — a requeued job starts again ***"), warn=True) + else: + _row("Resume", f"from the last checkpoint in {run_dir / 'checkpoints'}") + + _section("Output") + _row("Run directory", str(run_dir)) + for frozen, source, differs in LlamaFactoryBackend().kept_frozen_inputs(config, run_dir): + name = frozen.relative_to(run_dir.resolve()) + if differs: + text = f"*** {name} differs from {source} — the frozen copy will NOT be replaced ***" + else: + text = f"*** {name} already exists in the run directory — it will NOT be replaced ***" + _row("Frozen input", _yellow(text), warn=True) + _row("Logging", _lf_value(lf, "report_to")) + if config.debug.enabled: + _row("Debug mode", _red("*** ENABLED — output dir may be overwritten ***"), warn=True) + else: + _row("Debug mode", "disabled") + if tokenize_only: + _row( + "Tokenize only", + _yellow( + f"*** ENABLED — saves to {lf.get('tokenized_path')} and exits, no training ***" + ), + warn=True, + ) + + +# LlamaFactory 0.9.5 and transformers defaults for the keys that the review shows. +_LF_DEFAULTS = { + "template": "parsed from the tokenizer chat template", + "cutoff_len": 2048, + "finetuning_type": "lora", + "pref_loss": "sigmoid", + "pref_beta": 0.1, + "num_train_epochs": 3.0, + "learning_rate": 5e-5, + "lr_scheduler_type": "linear", + "per_device_train_batch_size": 8, + "gradient_accumulation_steps": 1, + "save_strategy": "steps", + "save_steps": 500, + # The default changes between transformers releases. + "report_to": "transformers default", +} + + +def _lf_value(lf: dict, key: str) -> str: + """The value of *key*, or the default marked as such when the block does not set it.""" + if key in lf: + return str(lf[key]) + return f"{_LF_DEFAULTS[key]} (default)" + + +def _confirm(total_gpus: int) -> None: # ------------------------------------------------------------------ # Final confirmation # ------------------------------------------------------------------ diff --git a/src/post_training/utils/prefetch.py b/src/post_training/utils/prefetch.py index 3097114..22e742e 100644 --- a/src/post_training/utils/prefetch.py +++ b/src/post_training/utils/prefetch.py @@ -9,17 +9,23 @@ :func:`datasets.load_dataset` so that the processed Arrow cache is populated (``~/.cache/huggingface/datasets/``), which is what ``HF_DATASETS_OFFLINE=1`` requires at runtime. + +For ``backend: llamafactory`` the assets come from the ``llamafactory`` block +instead, and are fetched with the same arguments that LlamaFactory uses. """ from __future__ import annotations import dataclasses +import json import logging from pathlib import Path from typing import TYPE_CHECKING from datasets import load_dataset -from huggingface_hub import snapshot_download +from huggingface_hub import hf_hub_download, snapshot_download + +from post_training.backend import llamafactory_dataset_names if TYPE_CHECKING: from post_training.config import DatasetEntry, PostTrainingConfig @@ -54,7 +60,9 @@ def _is_local(path: str) -> bool: return Path(path).exists() -def _prefetch_model(name_or_path: str, revision: str | None = None) -> str: +def _prefetch_model( + name_or_path: str, revision: str | None = None, cache_dir: str | None = None +) -> str: """Ensure *name_or_path* is cached locally and return its local directory.""" if _is_local(name_or_path): logger.info("Model '%s' is a local path, skipping download.", name_or_path) @@ -62,7 +70,9 @@ def _prefetch_model(name_or_path: str, revision: str | None = None) -> str: logger.info( "Downloading model '%s' (revision=%s) to HF cache...", name_or_path, revision or "main" ) - local_path = snapshot_download(repo_id=name_or_path, repo_type="model", revision=revision) + local_path = snapshot_download( + repo_id=name_or_path, repo_type="model", revision=revision, cache_dir=cache_dir + ) logger.info("Model '%s' cached at '%s'.", name_or_path, local_path) return local_path @@ -132,6 +142,9 @@ def prefetch_assets(config: PostTrainingConfig) -> PrefetchedPaths: """ logger.info("Pre-fetching assets for offline run...") + if config.backend == "llamafactory": + return _prefetch_llamafactory(config) + model_path = _prefetch_model(config.model.name_or_path, revision=config.model.revision) tokenizer_name_or_path, tokenizer_revision = config.model.resolve_tokenizer() @@ -153,3 +166,142 @@ def prefetch_assets(config: PostTrainingConfig) -> PrefetchedPaths: logger.info("All assets pre-fetched successfully.") return PrefetchedPaths(model=model_path, tokenizer=tokenizer_path, ref_model=ref_model_path) + + +# --------------------------------------------------------------------------- +# LlamaFactory +# --------------------------------------------------------------------------- + +_LLAMAFACTORY_REGISTRY = "dataset_info.json" + + +@dataclasses.dataclass +class _HubDataset: + """One load_dataset call, with the arguments LlamaFactory passes for a Hub dataset.""" + + name: str # the name used in `dataset` / `eval_dataset` + path: str + subset: str | None = None + folder: str | None = None + split: str = "train" + + +def _llamafactory_has_tokenized_data(path: str) -> bool: + """Same test as LlamaFactory's has_tokenized_data(): a directory that is not empty.""" + tokenized = Path(path) + return tokenized.is_dir() and any(tokenized.iterdir()) + + +def _llamafactory_source(entry: dict) -> str: + """The load_from that LlamaFactory's get_dataset_list() picks for a registry entry. + + LlamaFactory also lets USE_MODELSCOPE_HUB / USE_OPENMIND_HUB prefer those hubs; + the jobs here do not set them. + """ + has_hf, has_ms, has_om = ("hf_hub_url" in entry, "ms_hub_url" in entry, "om_hub_url" in entry) + if has_hf or has_ms or has_om: + if has_ms and not has_hf: + return "ms_hub" + if has_om and not has_hf: + return "om_hub" + return "hf_hub" + if "script_url" in entry: + return "script" + if "cloud_file_name" in entry: + return "cloud_file" + return "file" + + +def _read_llamafactory_registry(dataset_dir: str) -> dict: + if dataset_dir.startswith("REMOTE:"): + registry = hf_hub_download( + repo_id=dataset_dir[len("REMOTE:") :], + filename=_LLAMAFACTORY_REGISTRY, + repo_type="dataset", + ) + else: + registry = Path(dataset_dir) / _LLAMAFACTORY_REGISTRY + return json.loads(Path(registry).read_text()) + + +def _plan_llamafactory_datasets(lf: dict) -> tuple[list[_HubDataset], list[str]]: + """Return the Hub loads to pre-fetch, and the datasets the prefetch cannot fetch.""" + names = llamafactory_dataset_names(lf) + tokenized_path = lf.get("tokenized_path") + if isinstance(tokenized_path, str) and _llamafactory_has_tokenized_data(tokenized_path): + logger.info( + "tokenized_path '%s' holds data; LlamaFactory ignores the datasets.", tokenized_path + ) + return [], [] + if not names: + return [], [] + + dataset_dir = lf.get("dataset_dir") or "data" + loads: list[_HubDataset] = [] + unfetchable: list[str] = [] + if dataset_dir == "ONLINE": + loads = [_HubDataset(name=name, path=name) for name in names] + else: + registry = _read_llamafactory_registry(dataset_dir) + for name in names: + if name not in registry: + raise ValueError( + f"Dataset '{name}' is not defined in {dataset_dir}/{_LLAMAFACTORY_REGISTRY}." + ) + entry = registry[name] + source = _llamafactory_source(entry) + if source == "hf_hub": + loads.append( + _HubDataset( + name=name, + path=entry["hf_hub_url"], + subset=entry.get("subset"), + folder=entry.get("folder"), + split=entry.get("split", "train"), + ) + ) + elif source != "file": + unfetchable.append(f"{name} ({source})") + + # A streaming load reads the Hub directly and does not use the prepared cache. + if lf.get("streaming") and loads: + unfetchable.extend(f"{load.name} (streaming)" for load in loads) + loads = [] + return loads, unfetchable + + +def _prefetch_llamafactory(config: PostTrainingConfig) -> PrefetchedPaths: + """Download what llamafactory-cli loads: the model, ref_model and the Hub datasets. + + The dataset check runs first, so an offline run that could not load a dataset + is refused before any large download. + """ + lf = config.llamafactory + loads, unfetchable = _plan_llamafactory_datasets(lf) + if unfetchable: + if config.offline: + raise ValueError( + "offline=true, but the prefetch cannot download these datasets, so the job " + f"could not load them: {', '.join(unfetchable)}." + ) + logger.warning( + "Not pre-fetching these datasets; the job loads them itself: %s", + ", ".join(unfetchable), + ) + + # LlamaFactory loads ref_model with a copy of the policy's model arguments. + cache_dir = lf.get("cache_dir") + revision = lf.get("model_revision") + model_path = _prefetch_model(lf["model_name_or_path"], revision=revision, cache_dir=cache_dir) + ref_model_path = None + if lf.get("ref_model"): + ref_model_path = _prefetch_model(lf["ref_model"], revision=revision, cache_dir=cache_dir) + + for load in loads: + logger.info("Downloading dataset '%s' ('%s') to HF cache...", load.name, load.path) + load_dataset( + load.path, name=load.subset, data_dir=load.folder, split=load.split, cache_dir=cache_dir + ) + + logger.info("All assets pre-fetched successfully.") + return PrefetchedPaths(model=model_path, tokenizer=model_path, ref_model=ref_model_path) diff --git a/tests/test_backend.py b/tests/test_backend.py index eebe4d3..c51ad25 100644 --- a/tests/test_backend.py +++ b/tests/test_backend.py @@ -1,8 +1,14 @@ -"""Tests for TRLBackend.post_freeze() — the submission-time source freeze.""" +"""Tests for the submission-time behavior of TRLBackend and LlamaFactoryBackend.""" +import json import logging +from pathlib import Path -from post_training.backend import TRLBackend +import pytest +import yaml + +import post_training.backend as backend_module +from post_training.backend import LlamaFactoryBackend, TRLBackend from post_training.config import ContainerConfig, PostTrainingConfig @@ -123,3 +129,544 @@ def test_post_freeze_recovers_from_interrupted_copy(tmp_path, monkeypatch): assert (frozen_src / "marker.txt").read_text() == "live" assert not (frozen_src / "half-written.py").exists() assert not tmp_src.exists() + + +# --------------------------------------------------------------------------- +# LlamaFactoryBackend — output directory and resume +# --------------------------------------------------------------------------- + +_REPO_ROOT = Path(__file__).resolve().parent.parent + + +def _llamafactory_config(**block) -> PostTrainingConfig: + cfg = PostTrainingConfig() + cfg.backend = "llamafactory" + cfg.container = ContainerConfig( + image="/shared/containers/llamafactory.sif", + bind_mounts=["/scratch:/scratch"], + env_file="/shared/env/cluster.env", + ) + cfg.llamafactory = { + "model_name_or_path": "org/model", + "stage": "sft", + "do_train": True, + **block, + } + return cfg + + +def _frozen_llamafactory_config(run_dir) -> dict: + return yaml.safe_load((run_dir / "llamafactory_config.yaml").read_text()) + + +def test_llamafactory_validate_requires_model_name_or_path(): + cfg = _llamafactory_config() + del cfg.llamafactory["model_name_or_path"] + + with pytest.raises(ValueError, match="llamafactory.model_name_or_path"): + LlamaFactoryBackend().validate(cfg) + + +def test_llamafactory_validate_rejects_output_dir(): + """The submission owns output_dir, so a user value is refused before any queue wait.""" + with pytest.raises(ValueError, match="llamafactory.output_dir"): + LlamaFactoryBackend().validate(_llamafactory_config(output_dir="output/my-run")) + + +def test_llamafactory_validate_rejects_overwrite_output_dir(): + """overwrite_output_dir: true turns off LlamaFactory's resume after a requeue.""" + with pytest.raises(ValueError, match="llamafactory.overwrite_output_dir"): + LlamaFactoryBackend().validate(_llamafactory_config(overwrite_output_dir=True)) + + +def test_llamafactory_validate_allows_overwrite_output_dir_false(): + LlamaFactoryBackend().validate(_llamafactory_config(overwrite_output_dir=False)) + + +@pytest.mark.parametrize("stage", ["pt", "sft", "rm", "ppo", "dpo", "kto"]) +def test_llamafactory_validate_accepts_matching_stage_and_method(stage): + cfg = _llamafactory_config(stage=stage) + cfg.method = stage + + LlamaFactoryBackend().validate(cfg) + + +def test_llamafactory_validate_defaults_the_stage_to_sft(): + cfg = _llamafactory_config() + del cfg.llamafactory["stage"] + + LlamaFactoryBackend().validate(cfg) + + +@pytest.mark.parametrize("stage", ["orpo", "SFT", None]) +def test_llamafactory_validate_rejects_unknown_stage(stage): + cfg = _llamafactory_config(stage=stage) + + with pytest.raises(ValueError, match="llamafactory.stage .* Choose from: pt, sft"): + LlamaFactoryBackend().validate(cfg) + + +def test_llamafactory_validate_rejects_method_that_differs_from_stage(): + cfg = _llamafactory_config(stage="dpo") + cfg.method = "sft" + + with pytest.raises(ValueError, match="method 'sft' does not match llamafactory.stage 'dpo'"): + LlamaFactoryBackend().validate(cfg) + + +def test_llamafactory_validate_rejects_default_stage_with_other_method(): + """Without llamafactory.stage, LlamaFactory trains sft, so method must be sft.""" + cfg = _llamafactory_config() + del cfg.llamafactory["stage"] + cfg.method = "dpo" + + with pytest.raises(ValueError, match="does not match llamafactory.stage 'sft'"): + LlamaFactoryBackend().validate(cfg) + + +def test_llamafactory_example_config_is_valid(monkeypatch): + """The shipped example config passes validation from the repository root.""" + monkeypatch.chdir(_REPO_ROOT) + + PostTrainingConfig.load(_REPO_ROOT / "configs" / "llamafactory" / "long-context.yaml") + + +def test_llamafactory_run_dir_has_checkpoints_subdir(): + assert LlamaFactoryBackend().run_dir_subdirs() == ["checkpoints"] + + +def test_llamafactory_post_freeze_writes_resumable_output_dir(tmp_path, monkeypatch): + """Checkpoints go to /checkpoints with resume enabled; config.yaml keeps the user block.""" + monkeypatch.chdir(tmp_path) + cfg = _llamafactory_config(learning_rate=1e-5) + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + LlamaFactoryBackend().post_freeze(cfg, run_dir) + + frozen = _frozen_llamafactory_config(run_dir) + assert frozen["output_dir"] == str(run_dir.resolve() / "checkpoints") + assert frozen["overwrite_output_dir"] is False + assert frozen["learning_rate"] == 1e-5 + assert "output_dir" not in cfg.llamafactory + assert "overwrite_output_dir" not in cfg.llamafactory + + +# --------------------------------------------------------------------------- +# LlamaFactoryBackend — frozen DeepSpeed and dataset configs, input paths +# --------------------------------------------------------------------------- + +_DATASET_INFO = { + "local_file": {"file_name": "sft.jsonl", "formatting": "sharegpt"}, + "absolute_file": {"file_name": "/datasets/abs.jsonl"}, + "local_script": {"script_url": "my_script"}, + "hub": {"hf_hub_url": "org/dataset"}, +} + + +def _make_llamafactory_repo(repo_dir: Path) -> None: + """A submission directory with a DeepSpeed JSON and a LlamaFactory dataset registry.""" + (repo_dir / "configs" / "deepspeed").mkdir(parents=True) + (repo_dir / "configs" / "deepspeed" / "zero3.json").write_text( + '{"zero_optimization": {"stage": 3}}' + ) + (repo_dir / "data" / "llamafactory").mkdir(parents=True) + (repo_dir / "data" / "llamafactory" / "dataset_info.json").write_text(json.dumps(_DATASET_INFO)) + + +def test_llamafactory_validate_rejects_missing_deepspeed_file(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + + with pytest.raises(ValueError, match="llamafactory.deepspeed"): + LlamaFactoryBackend().validate(_llamafactory_config(deepspeed="configs/ds/missing.json")) + + +def test_llamafactory_validate_accepts_inline_deepspeed_dict(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + + LlamaFactoryBackend().validate( + _llamafactory_config(deepspeed={"zero_optimization": {"stage": 2}}) + ) + + +@pytest.mark.parametrize( + "block", + [ + {"dataset": "long_sft", "dataset_dir": "data/llamafactory"}, + {"dataset": "long_sft"}, # LlamaFactory's default dataset_dir is "data" + {"eval_dataset": "long_sft_eval"}, + ], +) +def test_llamafactory_validate_rejects_missing_dataset_info(tmp_path, monkeypatch, block): + """A named dataset needs /dataset_info.json, or the job stops at start.""" + monkeypatch.chdir(tmp_path) + + with pytest.raises(ValueError, match="dataset_info.json"): + LlamaFactoryBackend().validate(_llamafactory_config(**block)) + + +@pytest.mark.parametrize("dataset_dir", ["ONLINE", "REMOTE:org/registry"]) +def test_llamafactory_validate_skips_registry_check_for_hub_registries( + tmp_path, monkeypatch, dataset_dir +): + monkeypatch.chdir(tmp_path) + + LlamaFactoryBackend().validate( + _llamafactory_config(dataset="org/dataset", dataset_dir=dataset_dir) + ) + + +def test_llamafactory_post_freeze_copies_deepspeed_json(tmp_path, monkeypatch): + repo_dir = tmp_path / "repo" + _make_llamafactory_repo(repo_dir) + monkeypatch.chdir(repo_dir) + cfg = _llamafactory_config(deepspeed="configs/deepspeed/zero3.json") + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + LlamaFactoryBackend().post_freeze(cfg, run_dir) + + copy_path = run_dir.resolve() / "llamafactory" / "zero3.json" + assert _frozen_llamafactory_config(run_dir)["deepspeed"] == str(copy_path) + assert copy_path.read_text() == (repo_dir / "configs" / "deepspeed" / "zero3.json").read_text() + assert cfg.llamafactory["deepspeed"] == "configs/deepspeed/zero3.json" + + +def test_llamafactory_post_freeze_keeps_inline_deepspeed_dict(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + inline = {"zero_optimization": {"stage": 2}} + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + LlamaFactoryBackend().post_freeze(_llamafactory_config(deepspeed=inline), run_dir) + + assert _frozen_llamafactory_config(run_dir)["deepspeed"] == inline + + +def test_llamafactory_post_freeze_freezes_dataset_registry(tmp_path, monkeypatch): + """Relative local entries point back at the source; the job reads the frozen registry.""" + repo_dir = tmp_path / "repo" + _make_llamafactory_repo(repo_dir) + monkeypatch.chdir(repo_dir) + cfg = _llamafactory_config(dataset="local_file", dataset_dir="data/llamafactory") + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + LlamaFactoryBackend().post_freeze(cfg, run_dir) + + source_dir = (repo_dir / "data" / "llamafactory").resolve() + frozen_dir = run_dir.resolve() / "llamafactory" + frozen = _frozen_llamafactory_config(run_dir) + assert frozen["dataset_dir"] == str(frozen_dir) + assert frozen["media_dir"] == str(source_dir) + registry = json.loads((frozen_dir / "dataset_info.json").read_text()) + assert registry["local_file"] == { + "file_name": str(source_dir / "sft.jsonl"), + "formatting": "sharegpt", + } + assert registry["absolute_file"] == _DATASET_INFO["absolute_file"] + assert registry["local_script"] == {"script_url": str(source_dir / "my_script")} + assert registry["hub"] == _DATASET_INFO["hub"] + assert cfg.llamafactory["dataset_dir"] == "data/llamafactory" + + +def test_llamafactory_post_freeze_online_registry_is_not_copied(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + LlamaFactoryBackend().post_freeze( + _llamafactory_config(dataset="org/dataset", dataset_dir="ONLINE"), run_dir + ) + + frozen = _frozen_llamafactory_config(run_dir) + assert frozen["dataset_dir"] == "ONLINE" + assert "media_dir" not in frozen + assert not (run_dir / "llamafactory" / "dataset_info.json").exists() + + +def test_llamafactory_post_freeze_makes_existing_input_paths_absolute(tmp_path, monkeypatch): + """Existing relative input paths survive the job's cd; hub IDs, new paths and outputs do not change.""" + repo_dir = tmp_path / "repo" + for local in ("models/local-model", "adapters/a", "adapters/b", "media", "exports"): + (repo_dir / local).mkdir(parents=True) + monkeypatch.chdir(repo_dir) + cfg = _llamafactory_config( + model_name_or_path="models/local-model", + adapter_name_or_path="adapters/a,adapters/b", + media_dir="media", + ref_model="org/ref-model", + tokenized_path="data/tokenized", + export_dir="exports", + dataset_dir="ONLINE", + ) + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + LlamaFactoryBackend().post_freeze(cfg, run_dir) + + root = repo_dir.resolve() + frozen = _frozen_llamafactory_config(run_dir) + assert frozen["model_name_or_path"] == str(root / "models" / "local-model") + assert frozen["adapter_name_or_path"] == f"{root / 'adapters' / 'a'},{root / 'adapters' / 'b'}" + assert frozen["media_dir"] == str(root / "media") + assert frozen["ref_model"] == "org/ref-model" + assert frozen["tokenized_path"] == "data/tokenized" + assert frozen["export_dir"] == "exports" + + +def test_llamafactory_post_freeze_makes_v095_input_paths_absolute(tmp_path, monkeypatch): + """The KTransformers and HyperParallel input files that LlamaFactory 0.9.5 added.""" + repo_dir = tmp_path / "repo" + (repo_dir / "kt").mkdir(parents=True) + for local in ("kt/weights.kt", "kt/experts.safetensors", "hyper_parallel.json"): + (repo_dir / local).write_text("{}") + monkeypatch.chdir(repo_dir) + cfg = _llamafactory_config( + kt_weight_path="kt/weights.kt", + kt_expert_checkpoint_path="kt/experts.safetensors", + hyper_parallel_args="hyper_parallel.json", + profiler_output_dir="profiler", + dataset_dir="ONLINE", + ) + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + LlamaFactoryBackend().post_freeze(cfg, run_dir) + + root = repo_dir.resolve() + frozen = _frozen_llamafactory_config(run_dir) + assert frozen["kt_weight_path"] == str(root / "kt" / "weights.kt") + assert frozen["kt_expert_checkpoint_path"] == str(root / "kt" / "experts.safetensors") + assert frozen["hyper_parallel_args"] == str(root / "hyper_parallel.json") + assert frozen["profiler_output_dir"] == "profiler" + + +def test_llamafactory_post_freeze_keeps_existing_copies(tmp_path, monkeypatch, caplog): + """A new submit into the same run directory keeps the copies its checkpoints were made with.""" + repo_dir = tmp_path / "repo" + _make_llamafactory_repo(repo_dir) + monkeypatch.chdir(repo_dir) + cfg = _llamafactory_config( + deepspeed="configs/deepspeed/zero3.json", + dataset="local_file", + dataset_dir="data/llamafactory", + ) + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + LlamaFactoryBackend().post_freeze(cfg, run_dir) + frozen_dir = run_dir / "llamafactory" + first_deepspeed = (frozen_dir / "zero3.json").read_text() + first_registry = (frozen_dir / "dataset_info.json").read_text() + + (repo_dir / "configs" / "deepspeed" / "zero3.json").write_text( + '{"zero_optimization": {"stage": 2}}' + ) + (repo_dir / "data" / "llamafactory" / "dataset_info.json").write_text("{}") + with caplog.at_level(logging.WARNING): + LlamaFactoryBackend().post_freeze(cfg, run_dir) + + assert (frozen_dir / "zero3.json").read_text() == first_deepspeed + assert (frozen_dir / "dataset_info.json").read_text() == first_registry + assert f"{frozen_dir / 'zero3.json'} differs from configs/deepspeed/zero3.json" in caplog.text + assert "dataset_info.json differs from" in caplog.text + + +def test_llamafactory_post_freeze_says_when_a_kept_copy_is_unchanged(tmp_path, monkeypatch, caplog): + repo_dir = tmp_path / "repo" + _make_llamafactory_repo(repo_dir) + monkeypatch.chdir(repo_dir) + cfg = _llamafactory_config(deepspeed="configs/deepspeed/zero3.json", dataset_dir="ONLINE") + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + LlamaFactoryBackend().post_freeze(cfg, run_dir) + + with caplog.at_level(logging.WARNING): + LlamaFactoryBackend().post_freeze(cfg, run_dir) + + assert "zero3.json already exists — keeping the frozen copy." in caplog.text + assert "differs" not in caplog.text + + +def test_llamafactory_kept_frozen_inputs_marks_changed_sources(tmp_path, monkeypatch): + repo_dir = tmp_path / "repo" + _make_llamafactory_repo(repo_dir) + monkeypatch.chdir(repo_dir) + cfg = _llamafactory_config( + deepspeed="configs/deepspeed/zero3.json", + dataset="local_file", + dataset_dir="data/llamafactory", + ) + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + assert LlamaFactoryBackend().kept_frozen_inputs(cfg, run_dir) == [] + + LlamaFactoryBackend().post_freeze(cfg, run_dir) + (repo_dir / "data" / "llamafactory" / "dataset_info.json").write_text("{}") + + frozen_dir = run_dir.resolve() / "llamafactory" + root = repo_dir.resolve() + assert LlamaFactoryBackend().kept_frozen_inputs(cfg, run_dir) == [ + ( + frozen_dir / "llamafactory_tokenize.py", + backend_module._LLAMAFACTORY_TOKENIZE_SCRIPT, + False, + ), + (frozen_dir / "zero3.json", Path("configs/deepspeed/zero3.json"), False), + ( + frozen_dir / "dataset_info.json", + root / "data" / "llamafactory" / "dataset_info.json", + True, + ), + ] + + +def test_llamafactory_post_freeze_always_freezes_the_tokenize_script(tmp_path, monkeypatch): + """Every submission freezes the tokenize script; a later one keeps that copy.""" + + monkeypatch.chdir(tmp_path) + source = tmp_path / "llamafactory_tokenize.py" + source.write_text("# version 1\n") + monkeypatch.setattr(backend_module, "_LLAMAFACTORY_TOKENIZE_SCRIPT", source) + cfg = _llamafactory_config(dataset_dir="ONLINE") + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + frozen = run_dir.resolve() / "llamafactory" / "llamafactory_tokenize.py" + + LlamaFactoryBackend().post_freeze(cfg, run_dir) + source.write_text("# version 2\n") + LlamaFactoryBackend().post_freeze(cfg, run_dir) + + assert frozen.read_text() == "# version 1\n" + assert LlamaFactoryBackend().kept_frozen_inputs(cfg, run_dir) == [(frozen, source, True)] + + +def test_llamafactory_post_freeze_recovers_from_interrupted_copy(tmp_path, monkeypatch): + """A leftover .tmp from a killed submission is replaced by a complete copy.""" + repo_dir = tmp_path / "repo" + _make_llamafactory_repo(repo_dir) + monkeypatch.chdir(repo_dir) + run_dir = tmp_path / "outputs" / "my-run" + (run_dir / "llamafactory").mkdir(parents=True) + (run_dir / "llamafactory" / "zero3.json.tmp").write_text('{"zero_opt') + + LlamaFactoryBackend().post_freeze( + _llamafactory_config(deepspeed="configs/deepspeed/zero3.json"), run_dir + ) + + copy_path = run_dir / "llamafactory" / "zero3.json" + assert copy_path.read_text() == (repo_dir / "configs" / "deepspeed" / "zero3.json").read_text() + assert not (run_dir / "llamafactory" / "zero3.json.tmp").exists() + + +# --------------------------------------------------------------------------- +# Prefetched snapshots +# --------------------------------------------------------------------------- + + +def test_trl_post_freeze_accepts_prefetched_paths(tmp_path): + """TRL substitutes snapshots into config.yaml in submit.py, so post_freeze ignores them.""" + from post_training.utils.prefetch import PrefetchedPaths + + cfg = PostTrainingConfig() + cfg.container = None + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + TRLBackend().post_freeze(cfg, run_dir, PrefetchedPaths(model="/m", tokenizer="/m")) + + assert not (run_dir / "src").exists() + + +def test_llamafactory_post_freeze_writes_prefetched_snapshots(tmp_path, monkeypatch): + """The job loads the prefetched snapshot directories; config.yaml keeps the Hub IDs.""" + from post_training.utils.prefetch import PrefetchedPaths + + monkeypatch.chdir(tmp_path) + cfg = _llamafactory_config(model_name_or_path="org/model", ref_model="org/ref") + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + prefetched = PrefetchedPaths( + model="/cache/org/model@main", + tokenizer="/cache/org/model@main", + ref_model="/cache/org/ref@main", + ) + + LlamaFactoryBackend().post_freeze(cfg, run_dir, prefetched) + + frozen = _frozen_llamafactory_config(run_dir) + assert frozen["model_name_or_path"] == "/cache/org/model@main" + assert frozen["ref_model"] == "/cache/org/ref@main" + assert cfg.llamafactory["model_name_or_path"] == "org/model" + assert cfg.llamafactory["ref_model"] == "org/ref" + + +# --------------------------------------------------------------------------- +# LlamaFactoryBackend — run name from the llamafactory block +# --------------------------------------------------------------------------- + +_TIMESTAMP = "20260917-101500" + + +def _llamafactory_run_name(cfg: PostTrainingConfig) -> str: + return LlamaFactoryBackend().generate_run_name(cfg, _TIMESTAMP) + + +def test_llamafactory_run_name_single_dataset(): + """eval_dataset does not enter the name.""" + cfg = _llamafactory_config( + model_name_or_path="ontocord/1.7b-MixtureVitae-300BT", + dataset="long_sft", + eval_dataset="long_sft_eval", + ) + + assert ( + _llamafactory_run_name(cfg) + == f"llamafactory-sft-17b-mixturevitae-300bt-long_sft-{_TIMESTAMP}" + ) + + +@pytest.mark.parametrize( + "datasets", ["dolci_instruct_sft, Tulu3-SFT", ["dolci_instruct_sft", "Tulu3-SFT"]] +) +def test_llamafactory_run_name_joins_datasets_with_plus(datasets): + cfg = _llamafactory_config(dataset=datasets) + + assert ( + _llamafactory_run_name(cfg) + == f"llamafactory-sft-model-dolci_instruct_sft+tulu3-sft-{_TIMESTAMP}" + ) + + +def test_llamafactory_run_name_without_datasets(): + assert _llamafactory_run_name(_llamafactory_config()) == f"llamafactory-sft-model-{_TIMESTAMP}" + + +def test_llamafactory_run_name_uses_the_stage(): + cfg = _llamafactory_config(stage="dpo", dataset="prefs") + + assert _llamafactory_run_name(cfg) == f"llamafactory-dpo-model-prefs-{_TIMESTAMP}" + + +def test_llamafactory_run_name_ignores_the_top_level_method(): + """Without llamafactory.stage the name uses LlamaFactory's default stage, not the TRL key.""" + cfg = _llamafactory_config() + del cfg.llamafactory["stage"] + cfg.method = "dpo" + + assert _llamafactory_run_name(cfg) == f"llamafactory-sft-model-{_TIMESTAMP}" + + +def test_llamafactory_run_name_local_model_with_trailing_slash(): + cfg = _llamafactory_config(model_name_or_path="/scratch/models/Prelude-9B/") + + assert _llamafactory_run_name(cfg) == f"llamafactory-sft-prelude-9b-{_TIMESTAMP}" + + +def test_llamafactory_run_name_hub_dataset_id(): + cfg = _llamafactory_config(dataset_dir="ONLINE", dataset="allenai/Dolci-Instruct-SFT") + + assert ( + _llamafactory_run_name(cfg) + == f"llamafactory-sft-model-allenai_dolci-instruct-sft-{_TIMESTAMP}" + ) diff --git a/tests/test_config.py b/tests/test_config.py index ee30cbc..a541a2c 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -203,3 +203,123 @@ 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) + + +# --------------------------------------------------------------------------- +# backend: llamafactory takes all training settings from the llamafactory block +# --------------------------------------------------------------------------- + +_TRL_SECTION_NAMES = ( + "model", + "training", + "checkpointing", + "data", + "sft", + "dpo", + "deepspeed", + "accelerate", + "logging", +) + + +def _write_llamafactory_yaml(tmp_path, **extra): + config = { + "backend": "llamafactory", + "container": { + "image": "/shared/containers/llamafactory.sif", + "bind_mounts": ["/scratch:/scratch"], + "env_file": "/shared/env/cluster.env", + }, + "llamafactory": {"model_name_or_path": "org/model", "stage": "sft", "do_train": True}, + **extra, + } + path = tmp_path / "config.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def test_llamafactory_config_rejects_trl_sections(tmp_path): + path = _write_llamafactory_yaml( + tmp_path, training={"learning_rate": 1e-5, "seed": 1}, sft={"max_seq_length": 4096} + ) + + with pytest.raises(ValueError, match="training.learning_rate") as excinfo: + PostTrainingConfig.load(path) + assert "training.seed" in str(excinfo.value) + assert "sft.max_seq_length" in str(excinfo.value) + + +def test_llamafactory_config_rejects_trl_cli_override(tmp_path): + path = _write_llamafactory_yaml(tmp_path) + + with pytest.raises(ValueError, match="sft.packing"): + PostTrainingConfig.load(path, ["sft.packing=false"]) + + +def test_llamafactory_config_rejects_null_trl_section(tmp_path): + """A key written with a null value counts as set.""" + path = _write_llamafactory_yaml(tmp_path, deepspeed=None) + + with pytest.raises(ValueError, match="deepspeed"): + PostTrainingConfig.load(path) + + +def test_llamafactory_backend_from_cli_override_is_checked(tmp_path): + path = _write_llamafactory_yaml(tmp_path, backend="trl", training={"max_steps": 1}) + + with pytest.raises(ValueError, match="training.max_steps"): + PostTrainingConfig.load(path, ["backend=llamafactory"]) + + +def test_llamafactory_config_allows_top_level_method(tmp_path): + path = _write_llamafactory_yaml(tmp_path, method="sft") + + assert PostTrainingConfig.load(path).method == "sft" + + +def test_trl_config_with_all_sections_loads(tmp_path, monkeypatch): + path = tmp_path / "config.yaml" + path.write_text( + yaml.safe_dump( + { + "backend": "trl", + "container": None, + "model": {"name_or_path": "org/model"}, + "training": {"max_steps": 1}, + "checkpointing": {"save_steps": 10}, + "data": {"datasets": [{"name": "dummy", "path": "dummy/path"}]}, + "sft": {"max_seq_length": 128}, + "dpo": {"beta": 0.1}, + "deepspeed": None, + "accelerate": {"dynamo_backend": "no"}, + "logging": {"logging_steps": 5}, + } + ) + ) + + PostTrainingConfig.load(path) + + +def test_llamafactory_save_omits_trl_sections(tmp_path): + """The frozen config holds only the llamafactory block and the submission keys, and loads again.""" + config = PostTrainingConfig.load(_write_llamafactory_yaml(tmp_path)) + frozen_path = tmp_path / "frozen.yaml" + + config.save(frozen_path) + + frozen = yaml.safe_load(frozen_path.read_text()) + assert not set(frozen) & set(_TRL_SECTION_NAMES) + assert {"backend", "llamafactory", "container", "slurm", "paths"} <= set(frozen) + PostTrainingConfig.load(frozen_path) + + +def test_trl_save_keeps_all_sections(tmp_path): + config = PostTrainingConfig() + config.container = None + config.training.max_steps = 1 + frozen_path = tmp_path / "frozen.yaml" + + config.save(frozen_path) + + frozen = yaml.safe_load(frozen_path.read_text()) + assert set(_TRL_SECTION_NAMES) <= set(frozen) diff --git a/tests/test_container_job_script.py b/tests/test_container_job_script.py new file mode 100644 index 0000000..7026264 --- /dev/null +++ b/tests/test_container_job_script.py @@ -0,0 +1,213 @@ +"""Run the rendered container job scripts under bash with stub SLURM commands. + +The text checks in test_slurm_render.py cannot show what bash does with a +script under ``set -euo pipefail``. These tests execute the script instead. +""" + +import os +import shutil +import signal +import subprocess +import time +from pathlib import Path + +import pytest + +from post_training.config import ContainerConfig, PostTrainingConfig +from post_training.slurm.launcher import ( + REQUIRED_HF_CACHE_VARS, + render_llamafactory_slurm_script, + render_trl_container_slurm_script, +) + +BASH = shutil.which("bash") +pytestmark = pytest.mark.skipif(BASH is None, reason="bash is not available") + +# Stub bodies, keyed by command name. The stubs record their calls in $STUB_LOG_DIR. +_STUBS = { + "srun": ( + 'echo "$@" >> "$STUB_LOG_DIR/srun.log"\n' + 'if [ -n "${FAKE_SRUN_SLEEP:-}" ]; then\n' + ' touch "$STUB_LOG_DIR/srun.started"\n' + ' sleep "$FAKE_SRUN_SLEEP"\n' + "fi\n" + 'exit "${FAKE_SRUN_EXIT:-0}"\n' + ), + "scontrol": ( + 'case "$1" in\n' + ' show) echo "$SLURM_JOB_NODELIST" ;;\n' + " requeue)\n" + ' echo "$@" >> "$STUB_LOG_DIR/requeue.log"\n' + ' exit "${FAKE_REQUEUE_EXIT:-0}"\n' + " ;;\n" + "esac\n" + ), + "sbatch": 'echo "$@" >> "$STUB_LOG_DIR/sbatch.log"\n', + "getent": 'echo "127.0.0.1 $2"\n', +} + +_RENDERERS = { + "trl_container": lambda config, run_dir: render_trl_container_slurm_script( + config, run_dir, "configs/trl/sft.yaml" + ), + "llamafactory": render_llamafactory_slurm_script, +} + +_MAX_FAILURES = PostTrainingConfig().slurm.max_failures + + +@pytest.fixture(params=sorted(_RENDERERS)) +def job_script(request, tmp_path): + """Render one container job script that sources an empty cluster env file.""" + env_file = tmp_path / "cluster.env" + env_file.write_text("") + config = PostTrainingConfig() + config.container = ContainerConfig( + image="/shared/containers/image.sif", + bind_mounts=["/scratch:/scratch"], + env_file=str(env_file), + ) + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + return _RENDERERS[request.param](config, run_dir) + + +@pytest.fixture +def log_dir(tmp_path): + path = tmp_path / "stub_logs" + path.mkdir() + return path + + +@pytest.fixture +def job_env(tmp_path, log_dir): + """A fresh job environment: stubs first on PATH, SLURM job and HF cache variables.""" + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + for name, body in _STUBS.items(): + stub = bin_dir / name + stub.write_text("#!/bin/bash\n" + body) + stub.chmod(0o755) + return { + "PATH": f"{bin_dir}:{os.environ.get('PATH', '/usr/bin:/bin')}", + "STUB_LOG_DIR": str(log_dir), + "SLURM_JOB_ID": "12345", + "SLURM_JOB_NODELIST": "node001", + "SLURM_NNODES": "1", + "HF_HOME": "/hf", + "HF_HUB_CACHE": "/hf/hub", + "HUGGINGFACE_HUB_CACHE": "/hf/hub", + } + + +def _run_job(script: Path, env: dict[str, str]) -> subprocess.CompletedProcess: + return subprocess.run([BASH, str(script)], env=env, capture_output=True, text=True, timeout=30) + + +def _run_job_with_sigusr1( + script: Path, env: dict[str, str], log_dir: Path, tmp_path: Path +) -> tuple[int, str]: + """Run the job, send SIGUSR1 to the batch shell while srun runs, return (exit code, stderr). + + Output goes to files, not pipes: the sleeping srun stub outlives the batch + shell and would hold a pipe open. The stub is killed with its process group. + """ + env = {**env, "FAKE_SRUN_SLEEP": "5"} + stderr_path = tmp_path / "job.err" + with (tmp_path / "job.out").open("w") as stdout, stderr_path.open("w") as stderr: + proc = subprocess.Popen( + [BASH, str(script)], env=env, stdout=stdout, stderr=stderr, start_new_session=True + ) + try: + deadline = time.monotonic() + 10 + while not (log_dir / "srun.started").exists(): + assert time.monotonic() < deadline, "the srun stub did not start" + time.sleep(0.05) + proc.send_signal(signal.SIGUSR1) + returncode = proc.wait(timeout=30) + finally: + try: + os.killpg(proc.pid, signal.SIGKILL) + except ProcessLookupError: + pass + proc.wait() + return returncode, stderr_path.read_text() + + +# --------------------------------------------------------------------------- +# HF cache variables +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("missing", REQUIRED_HF_CACHE_VARS) +def test_unset_hf_variable_stops_before_srun(job_script, job_env, log_dir, missing): + """A missing HF cache variable stops the job with a clear message before srun starts.""" + del job_env[missing] + + result = _run_job(job_script, job_env) + + assert result.returncode == 1 + assert f"ERROR: {missing} not set" in result.stderr + assert "unbound variable" not in result.stderr + assert not (log_dir / "srun.log").exists() + assert not (log_dir / "sbatch.log").exists() + + +def test_hf_datasets_cache_is_optional(job_script, job_env, log_dir): + """An unset HF_DATASETS_CACHE falls back to $HF_HOME/datasets instead of stopping the job.""" + assert "HF_DATASETS_CACHE" not in job_env + + result = _run_job(job_script, job_env) + + assert result.returncode == 0, result.stderr + assert "HF_DATASETS_CACHE=/hf/datasets" in (log_dir / "srun.log").read_text() + + +# --------------------------------------------------------------------------- +# Failure handling and requeue +# --------------------------------------------------------------------------- + + +def test_failed_srun_is_counted_and_resubmitted(job_script, job_env, log_dir): + """A failed srun increments failure_count, resubmits the script and keeps the exit code.""" + job_env["FAKE_SRUN_EXIT"] = "3" + + result = _run_job(job_script, job_env) + + assert result.returncode == 3, result.stderr + assert (job_script.parent / "failure_count").read_text().strip() == "1" + assert (log_dir / "sbatch.log").read_text().strip() == str(job_script) + + +def test_failed_srun_at_max_failures_stops(job_script, job_env, log_dir): + """The failure that reaches max_failures stops the job instead of resubmitting.""" + (job_script.parent / "failure_count").write_text(f"{_MAX_FAILURES - 1}\n") + job_env["FAKE_SRUN_EXIT"] = "3" + + result = _run_job(job_script, job_env) + + assert result.returncode == 1, result.stderr + assert (job_script.parent / "failure_count").read_text().strip() == str(_MAX_FAILURES) + assert not (log_dir / "sbatch.log").exists() + + +def test_sigusr1_requeues_without_resubmit(job_script, job_env, log_dir, tmp_path): + """SIGUSR1 requeues the job, and the failure handler does not also resubmit it.""" + returncode, stderr = _run_job_with_sigusr1(job_script, job_env, log_dir, tmp_path) + + assert returncode == 0, stderr + assert (log_dir / "requeue.log").read_text().strip() == "requeue 12345" + assert not (log_dir / "sbatch.log").exists() + assert (job_script.parent / "failure_count").read_text().strip() == "0" + + +def test_sigusr1_with_failed_requeue_exits(job_script, job_env, log_dir, tmp_path): + """If scontrol requeue fails, the job exits at once with an error and does not resubmit.""" + job_env["FAKE_REQUEUE_EXIT"] = "1" + + returncode, stderr = _run_job_with_sigusr1(job_script, job_env, log_dir, tmp_path) + + assert returncode != 0 + assert "scontrol requeue failed" in stderr + assert not (log_dir / "sbatch.log").exists() + assert (job_script.parent / "failure_count").read_text().strip() == "0" diff --git a/tests/test_guardrails_llamafactory.py b/tests/test_guardrails_llamafactory.py new file mode 100644 index 0000000..5da53c3 --- /dev/null +++ b/tests/test_guardrails_llamafactory.py @@ -0,0 +1,131 @@ +"""Tests for the pre-submission review of LlamaFactory runs.""" + +import builtins + +import pytest + +from post_training.backend import LlamaFactoryBackend +from post_training.config import ContainerConfig, PostTrainingConfig +from post_training.utils.guardrails import run_guardrails + + +def _llamafactory_config(**block) -> PostTrainingConfig: + cfg = PostTrainingConfig() + cfg.backend = "llamafactory" + cfg.method = block.get("stage", "sft") + cfg.container = ContainerConfig( + image="/shared/containers/llamafactory.sif", + bind_mounts=["/scratch:/scratch"], + env_file="/shared/env/cluster.env", + ) + cfg.slurm.num_nodes = 2 + cfg.slurm.gpus_per_node = 4 + cfg.llamafactory = {"model_name_or_path": "org/lf-model", "dataset": "a,b", **block} + return cfg + + +def _review(cfg, tmp_path, monkeypatch, capsys) -> str: + monkeypatch.setattr(builtins, "input", lambda _: "y") + run_guardrails(cfg, tmp_path / "outputs" / "my-run") + return capsys.readouterr().out + + +def test_shows_the_llamafactory_block_not_trl_defaults(tmp_path, monkeypatch, capsys): + cfg = _llamafactory_config( + stage="sft", + finetuning_type="full", + template="qwen3", + per_device_train_batch_size=2, + gradient_accumulation_steps=3, + ) + + out = _review(cfg, tmp_path, monkeypatch, capsys) + + assert "org/lf-model" in out + assert "llamafactory 0.9.5" in out + assert "a, b" in out + assert "2 per device × 8 GPUs × 3 grad-acc = 48 effective" in out + # TRL defaults that have no effect on a LlamaFactory run. + assert PostTrainingConfig().model.name_or_path not in out + assert "Liger kernel" not in out + assert "Inference ckpts" not in out + + +def test_marks_unset_keys_with_their_defaults(tmp_path, monkeypatch, capsys): + out = _review(_llamafactory_config(), tmp_path, monkeypatch, capsys) + + assert "lora (default)" in out + assert "3.0 (default)" in out + assert "8 per device × 8 GPUs × 1 grad-acc = 64 effective" in out + + +def test_max_steps_replaces_epochs(tmp_path, monkeypatch, capsys): + out = _review(_llamafactory_config(max_steps=1000), tmp_path, monkeypatch, capsys) + + assert "1,000 steps" in out + assert "Epochs" not in out + + +@pytest.mark.parametrize("stage", ["rm", "ppo"]) +@pytest.mark.parametrize("finetuning_type", ["full", "freeze"]) +def test_warns_when_the_stage_cannot_resume(tmp_path, monkeypatch, capsys, stage, finetuning_type): + cfg = _llamafactory_config(stage=stage, finetuning_type=finetuning_type) + + out = _review(cfg, tmp_path, monkeypatch, capsys) + + assert f"cannot resume stage {stage} with {finetuning_type} tuning" in out + + +def test_warns_when_no_checkpoints_are_saved(tmp_path, monkeypatch, capsys): + out = _review(_llamafactory_config(save_strategy="no"), tmp_path, monkeypatch, capsys) + + assert "no checkpoints" in out + + +def test_resumes_from_run_dir_checkpoints(tmp_path, monkeypatch, capsys): + cfg = _llamafactory_config(stage="dpo", finetuning_type="full") + + out = _review(cfg, tmp_path, monkeypatch, capsys) + + assert f"from the last checkpoint in {tmp_path / 'outputs' / 'my-run' / 'checkpoints'}" in out + assert "Preference loss" in out + + +def _write_run_inputs(repo_dir, registry: str) -> None: + (repo_dir / "data").mkdir(parents=True, exist_ok=True) + (repo_dir / "data" / "dataset_info.json").write_text(registry) + + +def test_no_frozen_input_warning_for_a_new_run_directory(tmp_path, monkeypatch, capsys): + monkeypatch.chdir(tmp_path) + _write_run_inputs(tmp_path, '{"a": {"file_name": "a.json"}}') + + out = _review(_llamafactory_config(), tmp_path, monkeypatch, capsys) + + assert "Frozen input" not in out + + +def test_warns_about_a_kept_frozen_registry(tmp_path, monkeypatch, capsys): + monkeypatch.chdir(tmp_path) + _write_run_inputs(tmp_path, '{"a": {"file_name": "a.json"}}') + cfg = _llamafactory_config() + LlamaFactoryBackend().post_freeze(cfg, tmp_path / "outputs" / "my-run") + + out = _review(cfg, tmp_path, monkeypatch, capsys) + + assert "llamafactory/dataset_info.json already exists in the run directory" in out + assert "will NOT be replaced" in out + + +def test_warns_when_a_kept_frozen_registry_differs(tmp_path, monkeypatch, capsys): + monkeypatch.chdir(tmp_path) + _write_run_inputs(tmp_path, '{"a": {"file_name": "a.json"}}') + cfg = _llamafactory_config() + LlamaFactoryBackend().post_freeze(cfg, tmp_path / "outputs" / "my-run") + _write_run_inputs(tmp_path, '{"a": {"hf_hub_url": "org/a"}}') + + out = _review(cfg, tmp_path, monkeypatch, capsys) + + source = tmp_path.resolve() / "data" / "dataset_info.json" + assert f"llamafactory/dataset_info.json differs from {source}" in out + assert "the frozen copy will NOT be replaced" in out diff --git a/tests/test_llamafactory_tokenize.py b/tests/test_llamafactory_tokenize.py new file mode 100644 index 0000000..8f432c3 --- /dev/null +++ b/tests/test_llamafactory_tokenize.py @@ -0,0 +1,154 @@ +"""Tests for the --tokenize-only script of the LlamaFactory backend, without llamafactory.""" + +import importlib.util +import os +from pathlib import Path + +import pytest +import yaml + +# The script is not a post_training module: the container runs it from the run directory. +_SCRIPT = Path(__file__).resolve().parents[1] / "src/post_training/slurm/llamafactory_tokenize.py" +_spec = importlib.util.spec_from_file_location("llamafactory_tokenize", _SCRIPT) +_tokenize = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_tokenize) + + +def _write_config(tmp_path, tokenized_path) -> str: + config = tmp_path / "llamafactory_config.yaml" + config.write_text( + yaml.safe_dump( + { + "model_name_or_path": "org/model", + "stage": "sft", + "tokenized_path": str(tokenized_path), + "deepspeed": "/run/llamafactory/zero3.json", + } + ) + ) + return str(config) + + +@pytest.fixture +def tokenize_calls(monkeypatch): + """Replace the LlamaFactory step with one that saves a marker into tokenized_path.""" + calls = [] + + def fake_tokenize(args): + calls.append(args) + os.makedirs(args["tokenized_path"]) + Path(args["tokenized_path"], "dataset_dict.json").write_text("{}") + + monkeypatch.setattr(_tokenize, "tokenize", fake_tokenize) + return calls + + +def test_partial_path_is_named_after_the_slurm_job(monkeypatch): + monkeypatch.setenv("SLURM_JOB_ID", "4711") + + assert _tokenize.partial_path("/data/tokenized/sft/") == "/data/tokenized/sft.partial-4711" + + +def test_partial_path_uses_the_pid_outside_slurm(monkeypatch): + monkeypatch.delenv("SLURM_JOB_ID", raising=False) + + assert _tokenize.partial_path("/data/sft") == f"/data/sft.partial-{os.getpid()}" + + +def test_main_saves_into_its_partial_directory_and_publishes_it( + tmp_path, monkeypatch, tokenize_calls, capsys +): + monkeypatch.setenv("SLURM_JOB_ID", "4711") + tokenized = tmp_path / "tokenized" + + _tokenize.main(_write_config(tmp_path, tokenized)) + + assert [call["tokenized_path"] for call in tokenize_calls] == [f"{tokenized}.partial-4711"] + assert "deepspeed" not in tokenize_calls[0] + assert (tokenized / "dataset_dict.json").exists() + assert not Path(f"{tokenized}.partial-4711").exists() + assert f"Tokenized dataset saved at {tokenized}" in capsys.readouterr().out + + +def test_main_replaces_an_empty_tokenized_path_directory(tmp_path, monkeypatch, tokenize_calls): + monkeypatch.setenv("SLURM_JOB_ID", "4711") + tokenized = tmp_path / "tokenized" + tokenized.mkdir() + + _tokenize.main(_write_config(tmp_path, tokenized)) + + assert (tokenized / "dataset_dict.json").exists() + + +def test_main_removes_the_interrupted_save_of_a_previous_attempt( + tmp_path, monkeypatch, tokenize_calls +): + """A requeued job has the same job ID, so its partial directory must not be loaded.""" + monkeypatch.setenv("SLURM_JOB_ID", "4711") + tokenized = tmp_path / "tokenized" + stale = Path(f"{tokenized}.partial-4711") + stale.mkdir() + (stale / "state.json").write_text("interrupted") + + _tokenize.main(_write_config(tmp_path, tokenized)) + + assert sorted(p.name for p in tokenized.iterdir()) == ["dataset_dict.json"] + + +def test_main_does_nothing_when_tokenized_path_holds_data( + tmp_path, monkeypatch, tokenize_calls, capsys +): + tokenized = tmp_path / "tokenized" + tokenized.mkdir() + (tokenized / "dataset_dict.json").write_text("{}") + + _tokenize.main(_write_config(tmp_path, tokenized)) + + assert tokenize_calls == [] + assert "already holds data" in capsys.readouterr().out + + +def test_main_discards_its_data_when_another_job_published_first(tmp_path, monkeypatch, capsys): + """Two jobs on one tokenized_path: the second publish fails and leaves the first intact.""" + monkeypatch.setenv("SLURM_JOB_ID", "4711") + tokenized = tmp_path / "tokenized" + + def tokenize_while_another_job_publishes(args): + os.makedirs(args["tokenized_path"]) + Path(args["tokenized_path"], "dataset_dict.json").write_text("second job") + tokenized.mkdir() + (tokenized / "dataset_dict.json").write_text("first job") + + monkeypatch.setattr(_tokenize, "tokenize", tokenize_while_another_job_publishes) + + _tokenize.main(_write_config(tmp_path, tokenized)) + + assert (tokenized / "dataset_dict.json").read_text() == "first job" + assert not Path(f"{tokenized}.partial-4711").exists() + assert "Another job saved tokenized data" in capsys.readouterr().out + + +def test_publish_keeps_the_first_published_data(tmp_path): + tokenized = tmp_path / "tokenized" + first, second = tmp_path / "tokenized.partial-1", tmp_path / "tokenized.partial-2" + for partial in (first, second): + partial.mkdir() + (partial / "dataset_dict.json").write_text(partial.name) + + assert _tokenize.publish(str(first), str(tokenized)) is True + assert _tokenize.publish(str(second), str(tokenized)) is False + assert (tokenized / "dataset_dict.json").read_text() == "tokenized.partial-1" + assert not first.exists() + assert not second.exists() + + +def test_publish_raises_other_errors(tmp_path): + """A target that is a file is a setup error, not a lost race.""" + tokenized = tmp_path / "tokenized" + tokenized.write_text("not a directory") + partial = tmp_path / "tokenized.partial-1" + partial.mkdir() + + with pytest.raises(OSError): + _tokenize.publish(str(partial), str(tokenized)) + assert partial.exists() diff --git a/tests/test_prefetch.py b/tests/test_prefetch.py index a957c63..72bc699 100644 --- a/tests/test_prefetch.py +++ b/tests/test_prefetch.py @@ -1,5 +1,9 @@ """Tests for which repos prefetch_assets downloads.""" +import json +import logging +from pathlib import Path + import pytest from post_training.config import PostTrainingConfig @@ -11,7 +15,9 @@ def downloads(monkeypatch): """Record snapshot_download calls instead of contacting the Hub.""" calls = [] - def fake_snapshot_download(repo_id, repo_type=None, revision=None, ignore_patterns=None): + def fake_snapshot_download( + repo_id, repo_type=None, revision=None, ignore_patterns=None, cache_dir=None + ): calls.append((repo_id, revision, ignore_patterns)) return f"/cache/{repo_id}@{revision or 'main'}" @@ -78,3 +84,199 @@ def test_local_tokenizer_path_is_not_downloaded(downloads, tmp_path): assert downloads == [("org/model", None, None)] assert paths.tokenizer == str(tmp_path) + + +# --------------------------------------------------------------------------- +# LlamaFactory backend +# --------------------------------------------------------------------------- + + +class _FakeHub: + """Records the Hub calls that the LlamaFactory prefetch makes.""" + + def __init__(self, tmp_path): + self.tmp_path = tmp_path + self.models = [] + self.datasets = [] + self.registries = [] + self.remote = {} # repo id -> dataset_info.json content served by hf_hub_download + + def snapshot_download( + self, repo_id, repo_type=None, revision=None, ignore_patterns=None, cache_dir=None + ): + self.models.append((repo_id, revision, cache_dir)) + return f"/cache/{repo_id}@{revision or 'main'}" + + def load_dataset(self, path, **kwargs): + self.datasets.append({"path": path, **kwargs}) + + def hf_hub_download(self, repo_id, filename, repo_type=None, **kwargs): + self.registries.append((repo_id, filename, repo_type)) + dst = self.tmp_path / "remote" / repo_id / filename + dst.parent.mkdir(parents=True, exist_ok=True) + dst.write_text(json.dumps(self.remote[repo_id])) + return str(dst) + + +@pytest.fixture +def hub(monkeypatch, tmp_path): + fake = _FakeHub(tmp_path) + monkeypatch.setattr(prefetch, "snapshot_download", fake.snapshot_download) + monkeypatch.setattr(prefetch, "load_dataset", fake.load_dataset) + monkeypatch.setattr(prefetch, "hf_hub_download", fake.hf_hub_download, raising=False) + monkeypatch.chdir(tmp_path) + return fake + + +def _lf_config(offline: bool = False, **block) -> PostTrainingConfig: + config = PostTrainingConfig() + config.backend = "llamafactory" + config.offline = offline + config.llamafactory = {"model_name_or_path": "org/lf-model", "stage": "sft", **block} + return config + + +def _write_registry(entries: dict, dataset_dir: str = "data") -> None: + registry = Path(dataset_dir) / "dataset_info.json" + registry.parent.mkdir(parents=True, exist_ok=True) + registry.write_text(json.dumps(entries)) + + +def _hub_load(path, name=None, data_dir=None, split="train", cache_dir=None) -> dict: + return { + "path": path, + "name": name, + "data_dir": data_dir, + "split": split, + "cache_dir": cache_dir, + } + + +def test_llamafactory_downloads_its_model_not_the_trl_default(hub, tmp_path): + cache_dir = str(tmp_path / "hf-cache") + + paths = prefetch.prefetch_assets(_lf_config(model_revision="r1", cache_dir=cache_dir)) + + assert hub.models == [("org/lf-model", "r1", cache_dir)] + assert paths.model == "/cache/org/lf-model@r1" + assert paths.tokenizer == paths.model + assert paths.ref_model is None + + +def test_llamafactory_ref_model_uses_the_policy_revision(hub): + """LlamaFactory loads ref_model with the policy's model_revision, so the prefetch does too.""" + paths = prefetch.prefetch_assets(_lf_config(model_revision="r1", ref_model="org/ref")) + + assert hub.models == [("org/lf-model", "r1", None), ("org/ref", "r1", None)] + assert paths.ref_model == "/cache/org/ref@r1" + + +def test_llamafactory_local_model_is_not_downloaded(hub, tmp_path): + (tmp_path / "models" / "local").mkdir(parents=True) + + paths = prefetch.prefetch_assets(_lf_config(model_name_or_path="models/local")) + + assert hub.models == [] + assert paths.model == "models/local" + + +def test_llamafactory_loads_only_named_registry_entries(hub): + """Each named hub entry is loaded as LlamaFactory loads it; local files need no download.""" + _write_registry( + { + "sft": { + "hf_hub_url": "org/sft", + "subset": "chat", + "folder": "en", + "split": "validation", + }, + "unused": {"hf_hub_url": "org/unused"}, + "local": {"file_name": "local.jsonl"}, + "both": {"hf_hub_url": "org/both", "ms_hub_url": "ms/both"}, + "eval": {"hf_hub_url": "org/eval"}, + } + ) + + prefetch.prefetch_assets(_lf_config(dataset="sft, local,both", eval_dataset="eval")) + + assert hub.datasets == [ + _hub_load("org/sft", name="chat", data_dir="en", split="validation"), + _hub_load("org/both"), + _hub_load("org/eval"), + ] + + +def test_llamafactory_unknown_dataset_name_is_rejected(hub): + _write_registry({"sft": {"hf_hub_url": "org/sft"}}) + + with pytest.raises(ValueError, match="missing_set"): + prefetch.prefetch_assets(_lf_config(dataset="sft,missing_set")) + assert hub.models == [] + assert hub.datasets == [] + + +def test_llamafactory_online_names_load_with_split_train(hub): + prefetch.prefetch_assets(_lf_config(dataset_dir="ONLINE", dataset="org/a,org/b")) + + assert hub.datasets == [_hub_load("org/a"), _hub_load("org/b")] + + +def test_llamafactory_remote_registry_is_downloaded_first(hub): + hub.remote["org/registry"] = {"sft": {"hf_hub_url": "org/sft", "subset": "chat"}} + + prefetch.prefetch_assets(_lf_config(dataset_dir="REMOTE:org/registry", dataset="sft")) + + assert hub.registries == [("org/registry", "dataset_info.json", "dataset")] + assert hub.datasets == [_hub_load("org/sft", name="chat")] + + +_UNFETCHABLE_ENTRIES = [ + {"ms_hub_url": "ms/unfetchable"}, + {"om_hub_url": "om/unfetchable"}, + {"cloud_file_name": "s3://bucket/unfetchable.json"}, + {"script_url": "unfetchable_script"}, +] + + +@pytest.mark.parametrize("entry", _UNFETCHABLE_ENTRIES) +def test_llamafactory_offline_rejects_sources_it_cannot_fetch(hub, entry): + """An offline job could not load these, so submission stops before any download.""" + _write_registry({"sft": {"hf_hub_url": "org/sft"}, "unfetchable_set": entry}) + + with pytest.raises(ValueError, match="unfetchable_set"): + prefetch.prefetch_assets(_lf_config(offline=True, dataset="sft,unfetchable_set")) + assert hub.models == [] + assert hub.datasets == [] + + +def test_llamafactory_offline_rejects_streaming(hub): + _write_registry({"sft": {"hf_hub_url": "org/sft"}}) + + with pytest.raises(ValueError, match="streaming"): + prefetch.prefetch_assets(_lf_config(offline=True, dataset="sft", streaming=True)) + assert hub.models == [] + assert hub.datasets == [] + + +@pytest.mark.parametrize("entry", _UNFETCHABLE_ENTRIES) +def test_llamafactory_online_skips_sources_it_cannot_fetch(hub, caplog, entry): + """Without offline mode the job can load these itself, so the prefetch only logs them.""" + _write_registry({"sft": {"hf_hub_url": "org/sft"}, "unfetchable_set": entry}) + + with caplog.at_level(logging.INFO, logger=prefetch.logger.name): + prefetch.prefetch_assets(_lf_config(dataset="sft,unfetchable_set")) + + assert hub.models == [("org/lf-model", None, None)] + assert hub.datasets == [_hub_load("org/sft")] + assert "unfetchable_set" in caplog.text + + +def test_llamafactory_tokenized_data_skips_datasets(hub, tmp_path): + """LlamaFactory loads a non-empty tokenized_path and ignores the dataset arguments.""" + (tmp_path / "tokenized").mkdir() + (tmp_path / "tokenized" / "dataset_dict.json").write_text("{}") + + prefetch.prefetch_assets(_lf_config(dataset="not_in_any_registry", tokenized_path="tokenized")) + + assert hub.models == [("org/lf-model", None, None)] + assert hub.datasets == [] diff --git a/tests/test_slurm_render.py b/tests/test_slurm_render.py index a0da377..c42182a 100644 --- a/tests/test_slurm_render.py +++ b/tests/test_slurm_render.py @@ -1,5 +1,7 @@ """Automated tests for SLURM template rendering.""" +from pathlib import Path + import pytest from post_training.backend import LlamaFactoryBackend, TRLBackend @@ -140,7 +142,7 @@ def test_trl_container_qos_mem_absent_when_none(tmp_path, config): # --------------------------------------------------------------------------- -# environment modules — TRL container template +# environment modules — TRL container and LlamaFactory templates # --------------------------------------------------------------------------- @@ -189,6 +191,51 @@ def test_trl_container_host_setup_absent_when_unspecified(tmp_path, config): assert "module load" not in content +def test_llamafactory_modules_rendered(tmp_path, config): + """Configured host modules are loaded before the cluster environment file.""" + config.slurm.modules = ["singularity/4.1.0", "cuda/12.4"] + config.slurm.module_purge = True + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + content = render_llamafactory_slurm_script(config, run_dir).read_text() + + commands = [ + "module purge", + "module load singularity/4.1.0", + "module load cuda/12.4", + 'source "/shared/env/cluster.env"', + ] + assert all(command in content for command in commands) + assert [content.index(command) for command in commands] == sorted( + content.index(command) for command in commands + ) + + +def test_llamafactory_module_purge_omitted_when_disabled(tmp_path, config): + """Modules can be loaded without first purging the inherited module set.""" + config.slurm.modules = ["singularity/4.1.0"] + config.slurm.module_purge = False + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + content = render_llamafactory_slurm_script(config, run_dir).read_text() + + assert "module load singularity/4.1.0" in content + assert "module purge" not in content + + +def test_llamafactory_host_setup_absent_when_unspecified(tmp_path, config): + """No module setup is rendered when the module list is empty.""" + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + content = render_llamafactory_slurm_script(config, run_dir).read_text() + + assert "module purge" not in content + assert "module load" not in content + + # --------------------------------------------------------------------------- # --tokenize-only forwarding — TRL templates and backend dispatch # --------------------------------------------------------------------------- @@ -285,20 +332,37 @@ def test_trl_backend_forwards_tokenize_only_container(tmp_path, config): assert "--tokenize-only" in _train_invocation(script.read_text()) -def test_llamafactory_backend_ignores_tokenize_only(tmp_path, config): - """LlamaFactoryBackend accepts the kwarg silently and never emits --tokenize-only. - - submit.py raises before reaching this point, but the backend itself - must not crash if a caller passes the kwarg. - """ +def test_llamafactory_backend_tokenize_only_runs_the_tokenize_script(tmp_path, config): + """--tokenize-only runs the frozen tokenize script under torchrun, not llamafactory-cli.""" run_dir = tmp_path / "outputs" / "my-run" run_dir.mkdir(parents=True) script = LlamaFactoryBackend().render_slurm_script( config, run_dir, "configs/llamafactory/long-context.yaml", tokenize_only=True ) + content = script.read_text() + + tokenize_script = run_dir.resolve() / "llamafactory" / "llamafactory_tokenize.py" + lf_config = run_dir.resolve() / "llamafactory_config.yaml" + assert ( + f"PYTHONSAFEPATH=1 torchrun --standalone --nproc_per_node 1 {tokenize_script} {lf_config}" + in content + ) + assert "llamafactory-cli train" not in content + + +def test_llamafactory_backend_trains_by_default(tmp_path, config): + """Without --tokenize-only, the job trains and does not run the tokenize script.""" + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + script = LlamaFactoryBackend().render_slurm_script( + config, run_dir, "configs/llamafactory/long-context.yaml" + ) + content = script.read_text() - assert "--tokenize-only" not in script.read_text() + assert "llamafactory-cli train" in content + assert "llamafactory_tokenize.py" not in content # --------------------------------------------------------------------------- @@ -421,3 +485,95 @@ def test_llamafactory_container_custom_path(tmp_path, config): content = render_llamafactory_slurm_script(config, run_dir).read_text() assert 'export PATH=\\"/opt/conda/bin:/usr/local/bin:/usr/bin:/bin\\"' in content + + +# --------------------------------------------------------------------------- +# CUDA_DEVICE_MAX_CONNECTIONS — LlamaFactory template +# --------------------------------------------------------------------------- + + +def test_llamafactory_cuda_device_max_connections_not_set(tmp_path, config): + """The LlamaFactory job does not serialize CUDA streams via CUDA_DEVICE_MAX_CONNECTIONS.""" + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + content = render_llamafactory_slurm_script(config, run_dir).read_text() + + assert "export CUDA_DEVICE_MAX_CONNECTIONS" not in content + + +def test_llamafactory_forces_torchrun(tmp_path, config): + """One GPU per node still launches torchrun, so the nodes train as one job.""" + config.slurm.gpus_per_node = 1 + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + content = render_llamafactory_slurm_script(config, run_dir).read_text() + + assert "export FORCE_TORCHRUN=1" in content + + +# --------------------------------------------------------------------------- +# offline flag — LlamaFactory template +# --------------------------------------------------------------------------- + + +def test_llamafactory_offline_enabled(tmp_path, config): + """offline: true switches HF and WandB to offline mode inside the container.""" + config.offline = True + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + content = render_llamafactory_slurm_script(config, run_dir).read_text() + + assert "export HF_HUB_OFFLINE=1" in content + assert "export HF_DATASETS_OFFLINE=1" in content + assert "export TRANSFORMERS_OFFLINE=1" in content + assert "export WANDB_MODE=offline" in content + assert "OFFLINE=0" not in content + + +def test_llamafactory_offline_disabled(tmp_path, config): + """offline: false clears the HF offline flags, overriding any value from the env file.""" + config.offline = False + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + content = render_llamafactory_slurm_script(config, run_dir).read_text() + + assert "export HF_HUB_OFFLINE=0" in content + assert "export HF_DATASETS_OFFLINE=0" in content + assert "export TRANSFORMERS_OFFLINE=0" in content + assert "OFFLINE=1" not in content + assert "WANDB_MODE" not in content + + +# --------------------------------------------------------------------------- +# working directory and PYTHONPATH — LlamaFactory template +# --------------------------------------------------------------------------- + + +def test_llamafactory_runs_in_run_directory(tmp_path, config, monkeypatch): + """The job runs in the run directory, not in the directory that submit.py ran from.""" + submit_dir = tmp_path / "repo" + submit_dir.mkdir() + monkeypatch.chdir(submit_dir) + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + content = render_llamafactory_slurm_script(config, run_dir).read_text() + + assert f'cd \\"{run_dir.resolve()}\\"' in content + assert f"cd {Path.cwd()}" not in content + assert f"llamafactory-cli train {run_dir.resolve() / 'llamafactory_config.yaml'}" in content + + +def test_llamafactory_unsets_pythonpath(tmp_path, config): + """LlamaFactory does not import post_training, so no PYTHONPATH reaches the container.""" + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + content = render_llamafactory_slurm_script(config, run_dir).read_text() + + assert "unset PYTHONPATH" in content + assert "PYTHONPATH=" not in content diff --git a/tests/test_submit_hf_check.py b/tests/test_submit_hf_check.py new file mode 100644 index 0000000..bb8aaf2 --- /dev/null +++ b/tests/test_submit_hf_check.py @@ -0,0 +1,128 @@ +"""Tests for the submission-time HF cache variable check in scripts/submit.py. + +A container job without one of these variables stops as soon as it starts, +possibly after days in the queue, so submit.py must refuse it up front. +""" + +import importlib.util +import os +import sys +from pathlib import Path +from unittest import mock + +import pytest +import yaml + +import post_training.utils.prefetch as prefetch_module +from post_training.config import ContainerConfig, PostTrainingConfig +from post_training.slurm.launcher import REQUIRED_HF_CACHE_VARS + +# scripts/ is not a Python package, so load submit.py from its path. +_SUBMIT_PATH = Path(__file__).resolve().parent.parent / "scripts" / "submit.py" +_spec = importlib.util.spec_from_file_location("submit", _SUBMIT_PATH) +_submit = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_submit) + + +@pytest.fixture(autouse=True) +def clean_hf_env(): + """Hide the developer's HF cache variables and undo what _apply_hf_env_from_file sets.""" + with mock.patch.dict(os.environ): + for var in {*_submit._HF_CACHE_VARS, *REQUIRED_HF_CACHE_VARS}: + os.environ.pop(var, None) + yield + + +def _write_env_file(tmp_path: Path, skip: tuple[str, ...] = (), extra: str = "") -> Path: + """Write an env file that exports every required variable except those in *skip*.""" + lines = [ + f"export {var}=/cluster/hf/{var.lower()}\n" + for var in REQUIRED_HF_CACHE_VARS + if var not in skip + ] + env_file = tmp_path / "cluster.env" + env_file.write_text("".join(lines) + extra) + return env_file + + +def _container_config(env_file: Path) -> PostTrainingConfig: + config = PostTrainingConfig() + config.container = ContainerConfig( + image="/shared/containers/image.sif", + bind_mounts=["/scratch:/scratch"], + env_file=str(env_file), + ) + return config + + +def test_all_variables_exported_passes(tmp_path): + env_file = _write_env_file(tmp_path) + _submit._apply_hf_env_from_file(str(env_file)) + + _submit._check_required_hf_vars(_container_config(env_file)) + + +@pytest.mark.parametrize("missing", REQUIRED_HF_CACHE_VARS) +def test_missing_variable_is_rejected(tmp_path, missing): + env_file = _write_env_file(tmp_path, skip=(missing,)) + _submit._apply_hf_env_from_file(str(env_file)) + + with pytest.raises(ValueError, match=missing) as excinfo: + _submit._check_required_hf_vars(_container_config(env_file)) + assert str(env_file) in str(excinfo.value) + + +def test_variable_from_login_environment_passes(tmp_path, monkeypatch): + """sbatch passes the submitting shell's environment to the job, so a login-shell value counts.""" + env_file = _write_env_file(tmp_path, skip=("HF_HUB_CACHE",)) + monkeypatch.setenv("HF_HUB_CACHE", "/login/hf/hub") + _submit._apply_hf_env_from_file(str(env_file)) + + _submit._check_required_hf_vars(_container_config(env_file)) + + +def test_empty_variable_is_rejected(tmp_path): + """An empty value fails the job's [ -z ] test, so it is rejected like an unset one.""" + env_file = _write_env_file(tmp_path, skip=("HF_HUB_CACHE",), extra="export HF_HUB_CACHE=\n") + _submit._apply_hf_env_from_file(str(env_file)) + + with pytest.raises(ValueError, match="HF_HUB_CACHE"): + _submit._check_required_hf_vars(_container_config(env_file)) + + +def test_bare_metal_job_is_not_checked(): + """The bare-metal job script does not forward the HF cache variables.""" + config = PostTrainingConfig() + config.container = None + + _submit._check_required_hf_vars(config) + + +def test_submit_stops_before_prefetch_and_run_directory(tmp_path, monkeypatch): + """main() rejects the job before it downloads assets or creates the run directory.""" + env_file = _write_env_file(tmp_path, skip=("HF_HUB_CACHE",)) + output_base = tmp_path / "outputs" + config_path = tmp_path / "config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "training": {"max_steps": 1}, + "container": { + "image": "/shared/containers/image.sif", + "bind_mounts": ["/scratch:/scratch"], + "env_file": str(env_file), + }, + "paths": {"output_base": str(output_base)}, + } + ) + ) + + def fail_prefetch(config): + raise AssertionError("prefetch_assets must not run") + + monkeypatch.setattr(prefetch_module, "prefetch_assets", fail_prefetch) + monkeypatch.setattr(sys, "argv", ["submit.py", "--config", str(config_path), "--confirm"]) + + with pytest.raises(ValueError, match="HF_HUB_CACHE"): + _submit.main() + assert not output_base.exists() diff --git a/tests/test_submit_llamafactory.py b/tests/test_submit_llamafactory.py new file mode 100644 index 0000000..889190c --- /dev/null +++ b/tests/test_submit_llamafactory.py @@ -0,0 +1,176 @@ +"""End-to-end tests of scripts/submit.py for backend: llamafactory, without sbatch or the Hub.""" + +import importlib.util +import os +import sys +from pathlib import Path +from unittest import mock + +import pytest +import yaml + +import post_training.utils.prefetch as prefetch_module +from post_training.slurm.launcher import REQUIRED_HF_CACHE_VARS + +# scripts/ is not a Python package, so load submit.py from its path. +_SUBMIT_PATH = Path(__file__).resolve().parent.parent / "scripts" / "submit.py" +_spec = importlib.util.spec_from_file_location("submit", _SUBMIT_PATH) +_submit = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_submit) + +_SNAPSHOT = "/cache/org/lf-model@main" + + +@pytest.fixture(autouse=True) +def clean_hf_env(): + """Hide the developer's HF cache variables and undo what _apply_hf_env_from_file sets.""" + with mock.patch.dict(os.environ): + for var in {*_submit._HF_CACHE_VARS, *REQUIRED_HF_CACHE_VARS}: + os.environ.pop(var, None) + yield + + +@pytest.fixture +def submitted(tmp_path, monkeypatch): + """Run submit.py main() on a LlamaFactory config; return the run directory and job scripts.""" + env_file = tmp_path / "cluster.env" + env_file.write_text("".join(f"export {var}=/cluster/hf\n" for var in REQUIRED_HF_CACHE_VARS)) + config_path = tmp_path / "config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "backend": "llamafactory", + "run_name": "lf-run", + "container": { + "image": "/shared/containers/llamafactory.sif", + "bind_mounts": ["/scratch:/scratch"], + "env_file": str(env_file), + }, + "paths": {"output_base": str(tmp_path / "outputs")}, + "llamafactory": { + "model_name_or_path": "org/lf-model", + "stage": "sft", + "do_train": True, + }, + } + ) + ) + submitted_scripts = [] + + def fake_prefetch(config): + return prefetch_module.PrefetchedPaths(model=_SNAPSHOT, tokenizer=_SNAPSHOT) + + def fake_generate_and_submit(config, run_dir, config_path, *, tokenize_only=False): + submitted_scripts.append(run_dir) + return "12345" + + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(prefetch_module, "prefetch_assets", fake_prefetch) + monkeypatch.setattr(_submit, "generate_and_submit", fake_generate_and_submit) + monkeypatch.setattr(sys, "argv", ["submit.py", "--config", str(config_path), "--confirm"]) + + _submit.main() + + return tmp_path / "outputs" / "lf-run", submitted_scripts + + +def test_llamafactory_job_loads_the_prefetched_snapshot(submitted): + run_dir, submitted_scripts = submitted + + lf_config = yaml.safe_load((run_dir / "llamafactory_config.yaml").read_text()) + assert lf_config["model_name_or_path"] == _SNAPSHOT + assert len(submitted_scripts) == 1 + + +def test_llamafactory_config_yaml_keeps_the_hub_id_without_trl_sections(submitted): + """config.yaml keeps the user's llamafactory block and holds no TRL section.""" + run_dir, _ = submitted + + frozen = yaml.safe_load((run_dir / "config.yaml").read_text()) + assert frozen["llamafactory"]["model_name_or_path"] == "org/lf-model" + trl_sections = {"model", "training", "checkpointing", "data", "sft", "dpo", "deepspeed"} + assert not set(frozen) & (trl_sections | {"accelerate", "logging"}) + + +def test_llamafactory_resubmit_from_frozen_config(submitted, monkeypatch): + """A new submit from the frozen config.yaml passes the TRL-section check.""" + run_dir, submitted_scripts = submitted + frozen = str(run_dir / "config.yaml") + monkeypatch.setattr(sys, "argv", ["submit.py", "--config", frozen, "--confirm"]) + + _submit.main() + + assert len(submitted_scripts) == 2 + + +def _tokenize_only_submit(tmp_path, monkeypatch, llamafactory_block): + """Run submit.py --tokenize-only; return the run directory and the submit calls.""" + env_file = tmp_path / "cluster.env" + env_file.write_text("".join(f"export {var}=/cluster/hf\n" for var in REQUIRED_HF_CACHE_VARS)) + config_path = tmp_path / "config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "backend": "llamafactory", + "run_name": "lf-tok", + "container": { + "image": "/shared/containers/llamafactory.sif", + "bind_mounts": ["/scratch:/scratch"], + "env_file": str(env_file), + }, + "paths": {"output_base": str(tmp_path / "outputs")}, + "slurm": {"num_nodes": 4, "gpus_per_node": 8}, + "llamafactory": { + "model_name_or_path": "org/lf-model", + "stage": "sft", + "do_train": True, + **llamafactory_block, + }, + } + ) + ) + calls = [] + + def fake_generate_and_submit(config, run_dir, config_path, *, tokenize_only=False): + calls.append((config.slurm.num_nodes, config.slurm.gpus_per_node, tokenize_only)) + return "12345" + + monkeypatch.chdir(tmp_path) + monkeypatch.setattr( + prefetch_module, + "prefetch_assets", + lambda config: prefetch_module.PrefetchedPaths(model=_SNAPSHOT, tokenizer=_SNAPSHOT), + ) + monkeypatch.setattr(_submit, "generate_and_submit", fake_generate_and_submit) + monkeypatch.setattr( + sys, "argv", ["submit.py", "--config", str(config_path), "--confirm", "--tokenize-only"] + ) + + _submit.main() + + return tmp_path / "outputs" / "lf-tok", calls + + +def test_llamafactory_tokenize_only_makes_tokenized_path_absolute(tmp_path, monkeypatch): + """A relative tokenized_path starts at the submit directory, on 1 node with 1 GPU.""" + run_dir, calls = _tokenize_only_submit( + tmp_path, monkeypatch, {"tokenized_path": "data/tokenized/sft"} + ) + + lf_config = yaml.safe_load((run_dir / "llamafactory_config.yaml").read_text()) + assert lf_config["tokenized_path"] == str(tmp_path.resolve() / "data" / "tokenized" / "sft") + assert calls == [(1, 1, True)] + + +def test_llamafactory_tokenize_only_needs_tokenized_path(tmp_path, monkeypatch): + with pytest.raises(ValueError, match="needs llamafactory.tokenized_path"): + _tokenize_only_submit(tmp_path, monkeypatch, {}) + + +def test_llamafactory_tokenize_only_refuses_existing_data(tmp_path, monkeypatch): + tokenized = tmp_path / "data" / "tokenized" + tokenized.mkdir(parents=True) + (tokenized / "dataset_dict.json").write_text("{}") + + with pytest.raises(ValueError, match="already holds data"): + _tokenize_only_submit(tmp_path, monkeypatch, {"tokenized_path": "data/tokenized"})