feat: add reproducible reasoning SFT on LUMI

README.md

@@ -10,6 +10,7 @@ This repo supports two training backends:
 
 - [Quick Start](#quick-start)
 - [Run SFT using a pre-built container](#-run-sft-using-a-pre-built-container)
+- [Reasoning SFT on LUMI](docs/reasoning-sft.md)
 - [Project Structure](#-project-structure)
 - [Design Philosophy](#-design-philosophy)
 - [Feature Guide](#-feature-guide)
@@ -116,6 +117,12 @@ Tokenizing first keeps the multi-node allocation from sitting idle during CPU-bo
 
 The example config is [`configs/trl/prelude-sft.yaml`](configs/trl/prelude-sft.yaml). It fine-tunes a 9B checkpoint on LUMI, with the tokenizer from a separate repo. For your own run, copy it and replace the checkpoint, data, container paths, and SLURM account.
 
+For reasoning continuation, use the production-oriented
+[`configs/trl/reasoning-sft-lumi.yaml`](configs/trl/reasoning-sft-lumi.yaml) and
+the accompanying [reasoning SFT guide](docs/reasoning-sft.md). It adds a frozen
+FSDP launch profile, a pinned parent revision, 16K assistant-only training, and
+an explicit contract for a token-balanced materialized reasoning mixture.
+
 ### Before you start
 
 The login node needs only the base dependencies from [Installation](#installation), because the training stack lives in the container. Run every `submit.py` command from the repository root, inside that environment. Relative paths in the config (`container.env_file`, `paths.output_base`) resolve against the root, and `submit.py` copies the code from it.
@@ -305,6 +312,7 @@ data:
   datasets:
     - name: "my_dataset"
       path: "org/dataset"
+      revision: "<commit-sha>"  # pin Hub inputs for reproducible runs
       split: "train"
       weight: 1.0  # 1 = full dataset, <1 undersamples, >1 oversamples
 ```
@@ -399,7 +407,7 @@ You must specify exactly one determining factor for training duration in the `tr
 
 - **DeepSpeed**: configured inline under the top-level `deepspeed:` key (see the reference config below); set `deepspeed: null` to disable DeepSpeed entirely.
   To switch from ZeRO stage 2 to stage 3, bump `zero_optimization.stage` to `3` and add the `stage3_*` tuning keys — see `configs/deepspeed/zero3.yaml` for a full example.
-- **Accelerate flags**: the `accelerate` section in the YAML mirrors the CLI flags required for multi-node setups (`mixed_precision`, `dynamo_backend`, `rdzv_backend`, etc.).
+- **Accelerate flags**: the `accelerate` section in the YAML mirrors the CLI flags required for multi-node setups (`mixed_precision`, `dynamo_backend`, `rdzv_backend`, etc.). Set `accelerate.config_file` for a full profile such as FSDP. The renderer copies that file into the run's `slurm/` directory before submission, so a queued job cannot drift with the working checkout.
   These are used by the SLURM launcher to generate the correct job script.
 - **Self-healing**: the SLURM launcher (`src/post_training/slurm/`) supports auto-requeueing.
   - `slurm.signal_time_seconds` ensures the job saves a checkpoint and requeues itself before the wall time expires

configs/accelerate/fsdp-lumi.yaml

@@ -0,0 +1,14 @@
+# Multi-node full-shard profile for LUMI MI250X GCDs. The generated SLURM
+# launcher overrides process counts, machine rank, and rendezvous coordinates.
+compute_environment: LOCAL_MACHINE
+distributed_type: FSDP
+mixed_precision: bf16
+num_processes: 8
+fsdp_config:
+  fsdp_sharding_strategy: FULL_SHARD
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_use_orig_params: true
+  fsdp_offload_params: false
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_sync_module_states: true

configs/trl/prelude-sft.yaml

@@ -54,6 +54,7 @@ data:
   datasets:
     - name: "dolci-instruct-sft"
       path: "allenai/Dolci-Instruct-SFT"     # HuggingFace dataset path
+      revision: bd3c8f3a9b2cc5a9682e44b96ddd0bb2ff027221
       split: "train"
       weight: 1.0
       transform: null                        # null = already conversational

configs/trl/reasoning-sft-lumi.yaml

@@ -0,0 +1,117 @@
+# Reasoning continuation profile validated against the OpenEuroLLM 9B Qwen3
+# checkpoint family. Replace only the site-specific paths and the materialized
+# data artifact before running the tokenize-only gate described in the guide.
+method: sft
+backend: trl
+run_name: oellm-9b-anneal300b-reasoning-sft
+offline: true
+prefetch_assets: true
+
+container:
+  image: /scratch/<project>/containers/post-training-rocm7.2.4-py3.12-torch2.9.1-trl1.7.0-olmo-patched.sif
+  bind_mounts:
+    - /scratch/<project>
+  path: /opt/venv/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
+  env_file: env/lumi.env
+
+model:
+  name_or_path: Neonkraft/oellm-9b-256k-theta64m-prelude-anneal300b-instruct-sft
+  revision: 85bf18fb4f0bee6ac6270f06b1d1c6b3be200f31
+  attn_implementation: flash_attention_2
+  dtype: bfloat16
+  tokenizer_name_or_path: null
+  tokenizer_revision: null
+
+training:
+  # 64 packed 16K sequences/step * 500 steps = 524,288,000 training tokens.
+  num_training_tokens: 524288000
+  learning_rate: 1.5e-6
+  effective_batch_size: 64
+  per_device_train_batch_size: 1
+  warmup_steps: 15
+  lr_scheduler_type: cosine
+  lr_scheduler_kwargs: null
+  adam_beta1: 0.9
+  adam_beta2: 0.95
+  weight_decay: 0.0
+  adam_epsilon: 1.0e-8
+  gradient_checkpointing: true
+  gradient_checkpointing_kwargs:
+    use_reentrant: false
+  bf16: true
+  seed: 20260917
+  use_liger_kernel: true
+  liger_kernel_config:
+    fused_linear_cross_entropy: true
+    cross_entropy: false
+    rope: false
+    rms_norm: false
+    swiglu: false
+
+sft:
+  max_seq_length: 16384
+  packing: true
+  dataset_num_proc: 6
+  # Never train a reasoning trace cut before </think> or its final answer.
+  truncated_span_action: drop
+
+checkpointing:
+  save_strategy: steps
+  save_steps: 100
+  save_total_limit: 5
+  inference_checkpoint_steps: 100
+  inference_checkpoint_path: inference_checkpoints
+
+data:
+  chat_template: qwen3
+  num_proc: 6
+  seed: 20260917
+  datasets:
+    # The production recipe is token-balanced, globally deduplicated and quality
+    # gated before training. Point this at that immutable, materialized artifact;
+    # do not replace it with row weights over the raw source datasets.
+    - name: reasoning-anneal300b-v1
+      path: /scratch/<project>/<user>/post-training-data/reasoning-anneal300b-v1/train.parquet
+      split: train
+      weight: 1.0
+      transform: null
+
+# FSDP is the tested 8-node path for this 9B/16K profile. Keep DeepSpeed off so
+# Accelerate uses the frozen profile below.
+deepspeed: null
+
+accelerate:
+  config_file: configs/accelerate/fsdp-lumi.yaml
+  mixed_precision: bf16
+  use_deepspeed: false
+  same_network: true
+  rdzv_backend: static
+  dynamo_backend: "no"
+
+logging:
+  report_to: []
+  wandb_project: oellm-reasoning
+  logging_steps: 10
+  include_num_input_tokens_seen: non_padding
+
+slurm:
+  account: <project>
+  partition: standard-g
+  num_nodes: 8
+  gpus_per_node: 8
+  cpus_per_task: 56
+  mem: 256G
+  wall_time: "24:00:00"
+  job_name: oellm-reasoning-sft
+  signal_time_seconds: 300
+  max_failures: 1
+  modules: []
+  module_purge: false
+
+debug:
+  enabled: false
+  override_existing: false
+
+paths:
+  output_base: outputs
+  debug_base: outputs/debug

configs/trl/sft.yaml

@@ -89,6 +89,7 @@ data:
   datasets:
     - name: "nemotron_pt_v2"
       path: "nvidia/Nemotron-Post-Training-Dataset-v2"
+      revision: null                         # pin production runs to a Hub commit
       split: "stem"
       weight: 1.0                            # 1 = full dataset, <1 undersamples, >1 oversamples
       transform: null                        # null = already conversational
@@ -156,4 +157,3 @@ paths:
   output_base: "outputs"
   debug_base: "outputs/debug"
 
-

docs/reasoning-sft.md

@@ -0,0 +1,121 @@
+# Reasoning SFT on LUMI
+
+Reasoning continuation is supervised fine-tuning: the trainer does not need a
+separate `reasoning` method. It does need stricter data preparation and launch
+discipline than a general instruction run. The reference profile is
+[`configs/trl/reasoning-sft-lumi.yaml`](../configs/trl/reasoning-sft-lumi.yaml).
+
+## What the framework already enforces
+
+The TRL path is suitable for a production reasoning run because it:
+
+- applies the Qwen3 chat template with generation markers and refuses to start
+  if assistant-only loss cannot be constructed;
+- drops rows whose assistant answer contributes no supervised tokens;
+- can drop a row when the 16K boundary cuts through its supervised span, which
+  prevents training on an unterminated `<think>` block or final answer;
+- packs complete conversational records and derives the step count from a
+  training-token budget;
+- aligns generation EOS with the chat template's `<|im_end|>` turn terminator;
+- pins the parent model, tokenizer, and Hub datasets to revisions;
+- freezes the run config, source, Accelerate profile, and generated SLURM job;
+- supports a one-GPU `--tokenize-only` gate before the full allocation; and
+- resumes full checkpoints and writes lightweight inference checkpoints.
+
+The reference profile adds the multi-node FSDP configuration used by the 9B,
+16K LUMI run. The generated job still supplies the live node count, process
+count, machine rank and rendezvous address, so the same profile works for the
+one-GPU gate and the eight-node run.
+
+## Data contract
+
+The training input is one immutable, materialized Parquet dataset with a
+`messages` column. Each value is a list of `{role, content}` messages and the
+last turn is the assistant answer. Reasoning examples contain a non-empty
+`<think>...</think>` block followed by a non-empty final answer. Replay examples
+may contain a normal answer without a think block.
+
+Prepare the artifact before this trainer sees it. A production builder should:
+
+1. pin every source to a commit revision;
+2. normalize records to the `messages` contract without flattening prompt and
+   answer into one string;
+3. keep only complete reasoning traces and final answers;
+4. reject overlength rows rather than truncate them;
+5. reject exact lexical loops (the current run uses a repeated 30-gram at least
+   20 times as the strict threshold);
+6. deduplicate normalized user prompts across sources, in a deterministic
+   priority order;
+7. allocate mixture shares by rendered tokens, not rows; and
+8. write a manifest with source revisions, filter counts, selected row/token
+   counts, tokenizer/template identity, seed, and output SHA-256.
+
+The framework's ordinary `data.datasets[].weight` is deliberately a row-count
+multiplier. It is useful for similarly distributed instruction datasets, but
+it does **not** reproduce a token-balanced reasoning mixture whose traces range
+from hundreds to tens of thousands of tokens. The reference config therefore
+consumes the already-materialized mixture as one dataset.
+
+For the current conservative continuation recipe, the materialized artifact is
+524,288,000 rendered tokens: a 65% reasoning / 35% instruction-replay target,
+plus a fixed multilingual reasoning coverage floor consumed once before the
+weighted allocation. The exact source recipe and its manifest belong with the
+data artifact; the training repository records the immutable artifact used by
+the run.
+
+## Configure the run
+
+Copy the reference YAML and replace these site-specific values:
+
+- `container.image` and `container.bind_mounts`;
+- `container.env_file` with the shared Hugging Face cache variables;
+- `data.datasets[0].path` with the immutable Parquet artifact;
+- `slurm.account`; and
+- `paths.output_base` if the default is not on a persistent filesystem.
+
+Do not change the pinned model revision, Qwen3 template, 16K sequence length,
+assistant-only loss behavior, or Liger fused-loss settings without a new smoke
+test. The parent checkpoint contains its tokenizer, so the tokenizer follows
+the same pinned revision.
+
+## Gate, train, verify
+
+First render and tokenize through the same code path as production:
+
+```bash
+python scripts/submit.py \
+  --config configs/trl/reasoning-sft-lumi.yaml \
+  --tokenize-only
+```
+
+Read the generated `slurm-<job-id>.out` and `.err`. The gate passes only when:
+
+- the decoded preview uses Qwen3 ChatML and has a supervised assistant answer;
+- no row is cut part-way through its supervised span;
+- the resolved step count is 500;
+- the parent and dataset paths resolve from the shared cache/filesystem; and
+- trainer construction completes without an FSDP, FlashAttention or Liger
+  error.
+
+Then submit the unchanged configuration:
+
+```bash
+python scripts/submit.py --config configs/trl/reasoning-sft-lumi.yaml
+```
+
+After training starts, confirm that the full run reuses the tokenization cache,
+that loss is finite and decreasing, that every rank advances together, and
+that a resumable checkpoint is written at step 100. A release candidate should
+be compared with the parent on the same multilingual instruction, math/code
+reasoning, long-context, and repetition-loop probes. Publish the parent commit,
+data-manifest checksum, frozen config, environment, SLURM job IDs, selected
+checkpoint, and evaluation results in the model card.
+
+## Scope still outside this repository
+
+This trainer does not currently build token-balanced, cross-source-deduplicated
+mixtures, perform benchmark decontamination, or run the evaluation suite. Those
+steps should stay fail-closed in the data/evaluation pipelines and hand this
+repository immutable artifacts plus manifests. Keeping that boundary explicit
+prevents a convenient row-weighted training config from being mistaken for the
+audited production data recipe.

scripts/data.py

@@ -128,6 +128,8 @@ def _run_inspect(args: argparse.Namespace, cli_overrides: list[str]) -> None:
             load_kwargs["name"] = entry.subset
         if getattr(entry, "data_dir", None) is not None:
             load_kwargs["data_dir"] = entry.data_dir
+        if getattr(entry, "revision", None) is not None and not Path(entry.path).exists():
+            load_kwargs["revision"] = entry.revision
 
         ds = load_dataset(entry.path, split=entry.split, **load_kwargs)
 

src/post_training/backend.py

@@ -71,11 +71,26 @@ def _shorten_dataset_name(name: str) -> str:
 
 
 def _dataset_mix_hash(datasets: list) -> str:
+    # Every field that can change the realised rows belongs in the identity.
+    # Omitting (for example) the Hub revision can make two materially different
+    # runs resolve to the same output directory and resume each other's state.
     canonical = sorted(
-        [{"name": d.name, "path": d.path, "weight": d.weight} for d in datasets],
+        [
+            {
+                "name": d.name,
+                "path": d.path,
+                "revision": d.revision,
+                "data_dir": d.data_dir,
+                "subset": d.subset,
+                "split": d.split,
+                "weight": d.weight,
+                "transform": d.transform,
+            }
+            for d in datasets
+        ],
         key=lambda x: x["name"],
     )
-    return hashlib.sha256(json.dumps(canonical).encode()).hexdigest()[:8]
+    return hashlib.sha256(json.dumps(canonical, sort_keys=True).encode()).hexdigest()[:8]
 
 
 class TRLBackend(Backend):
@@ -168,7 +183,7 @@ def generate_run_name(self, config: PostTrainingConfig, timestamp: str) -> str:
         model_short = _shorten_model_name(config.model.name_or_path)
         datasets = config.data.datasets
         if len(datasets) == 1:
-            ds_part = _shorten_dataset_name(datasets[0].name)
+            ds_part = f"{_shorten_dataset_name(datasets[0].name)}_{_dataset_mix_hash(datasets)}"
         else:
             ds_part = f"mix_{_dataset_mix_hash(datasets)}"
         return f"{config.method}-{model_short}-{ds_part}-{timestamp}"

src/post_training/config.py

@@ -178,6 +178,8 @@ class DatasetEntry:
 
     name: str = MISSING
     path: str = MISSING
+    # Hub branch, tag, or commit. Ignored for existing local paths.
+    revision: str | None = None
     data_dir: str | None = None
     subset: str | None = None
     split: str = "train"
@@ -199,6 +201,10 @@ class DataConfig:
 class AccelerateConfig:
     """Flags forwarded to ``accelerate launch`` for explicit multi-node control."""
 
+    # Optional Accelerate YAML for launchers that are not expressible by the
+    # simple flags below (for example FSDP on LUMI). The SLURM renderer freezes
+    # a copy beside job.sh before submission.
+    config_file: str | None = None
     mixed_precision: str = "bf16"
     # Only takes effect when the top-level `deepspeed:` config is also set;
     # `deepspeed: null` disables DeepSpeed at launch regardless of this flag.