feat: add llamafactory

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 `<run_dir>/config.yaml` holds only the `llamafactory` block and the submission keys, so you can submit again from it.
+- The run name is `llamafactory-<stage>-<model>-<datasets>-<timestamp>`, 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 `<run_dir>/llamafactory/`, and the job runs in `<run_dir>`. 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 `<run_dir>`.
+- LlamaFactory writes checkpoints and the final model to `<run_dir>/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:

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
 

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 <think></think> 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

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"
+    }
   }
 }

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