Auto heal failing node

.gitignore

@@ -232,3 +232,4 @@ slurm_jupiter.sh
 slurm_juwels.sh
 logs/
 output/
+dump/

.gitmodules

@@ -4,3 +4,9 @@
 [submodule "submodules/titan-oellm"]
 	path = submodules/titan-oellm
 	url = git@github.com:OpenEuroLLM/titan-oellm.git
+[submodule "submodules/oellm_evals"]
+	path = submodules/oellm_evals
+	url = git@github.com:OpenEuroLLM/oellm-evals.git
+[submodule "submodules/Megatron-Bridge"]
+	path = submodules/Megatron-Bridge
+	url = git@github.com:NVIDIA-NeMo/Megatron-Bridge.git

README.md

@@ -95,6 +95,78 @@ python scripts/run_autoexp.py --config-name experiments/default
 
 ```
 
+### Training → conversion → evaluation chain (Megatron-Bridge + oellm-evals)
+
+A reference chain that trains in Megatron-LM, converts the
+torch_dist checkpoint to HuggingFace by way of Megatron-Bridge, and runs
+lm-eval-harness through `oellm-eval schedule` lives under
+`config/experiments/korbi/chain_qwen3_bridge_train_eval_<cluster>.yaml`
+for jupiter, juwels, leonardo, and lumi.
+
+The two new backends are:
+
+- **`megatron_bridge`** — `MegatronBridgeBackend`. Runs
+  `python -m oellm_autoexp.backends.megatron_bridge.run_export …`
+  inside the training container; reuses the SIF, no separate bridge
+  image is needed thanks to runtime shims in `run_export.py`
+  (`nvidia_resiliency_ext.__version__`,
+  `_clean_metadata_for_serialization`, a `LayerWiseDistributedOptimizer`
+  placeholder) plus the `container/megatron/patch_bridge_lazy_imports.py`
+  source patch applied once per checkout.
+- **`oellm_eval`** — `OELLMEvalBackend`. Wraps `oellm-eval schedule
+  --local true …` so the eval runs inside the SLURM job allocated by
+  oellm-autoexp (no extra sbatch from oellm-evals). Reads task definitions
+  from `submodules/oellm_evals/oellm/resources/task-groups.yaml`.
+
+Once per checkout, on a login node with internet, run the installer:
+
+```bash
+# Auto-detects the cluster from hostname; pass --cluster NAME to override.
+# Adds --prefetch to also pre-cache the eval datasets into HF_HOME.
+bash scripts/install_eval_env.sh --prefetch
+```
+
+That single script:
+
+- builds the right Python environment for the cluster — `uv venv` on
+  juwels/jupiter, `pip install --user` inside the eval container on
+  leonardo, `pip install --target` inside the rocm container on lumi
+- installs `oellm-evals` with the `[eval]` / `[eval-base]` extras from
+  `submodules/oellm_evals/pyproject.toml` (single source of truth for the
+  lm-eval dep set), plus `oellm-autoexp` and the `compoconf==0.1.14` pin
+- applies `container/megatron/patch_bridge_lazy_imports.py` so
+  `from megatron.bridge import AutoBridge` is tolerant of missing model
+  bridges
+- pulls the Qwen3-0.6B tokenizer's heavy files into
+  `oellm_autoexp/postprocess/resources/megatron_bridge/`
+- with `--prefetch`, calls `scripts/prefetch_datasets.py` inside the
+  matching env so HF datasets land in the cache layout the eval stage
+  will actually read from
+
+Set `HF_HOME` to the cluster-appropriate location before running with
+`--prefetch` — on Leonardo this is the dotted `.cache/huggingface` dir
+(because oellm-evals's eval script hardcodes
+`HF_DATASETS_CACHE=$HF_HOME/datasets` and Leonardo's `lm_eval` reads from
+that legacy layout); see
+[`docs/bridge_eval_setup.md`](docs/bridge_eval_setup.md) for the per-cluster
+values.
+
+Then submit a chain experiment as usual:
+
+```bash
+PYTHONPATH=. python scripts/run_autoexp.py \
+    --config-name experiments/korbi/chain_qwen3_bridge_train_eval_<cluster>
+```
+
+The orchestrator submits train, waits for `latest_checkpointed_iteration.txt`
+under the train output, submits convert, then submits eval once the HF
+`model.safetensors` appears. Each stage's logs land under
+`$OUTPUT_DIR/chain_qwen3_bridge_train_eval_<cluster>_{0,1,2}/`; per-task
+eval results are written to `…_2/eval/<timestamp>/results/<hash>_<ts>.json`.
+
+Per-cluster env / venv / container prerequisites for both backends are in
+[`docs/bridge_eval_setup.md`](docs/bridge_eval_setup.md).
+
 ### Monitoring sessions
 - Every submission drops `<monitoring_state_dir>/<session_id>/<job_id>.json`. Resuming is symmetric:
 ```bash

config/backend/megatron_bridge.yaml

@@ -0,0 +1,17 @@
+defaults:
+  - megatron_bridge: base
+  - _self_
+
+class_name: MegatronBridgeBackend
+
+# Where the Megatron-Bridge submodule lives. Override per-cluster if needed.
+bridge_root: ${oc.env:MEGATRON_BRIDGE_ROOT,submodules/Megatron-Bridge}
+
+# Where the bundled configs / templates / tokenizers live.
+resources: oellm_autoexp/postprocess/resources/megatron_bridge
+
+python_cmd: "python"
+keep_staging: false
+
+env:
+  HF_HOME: ${oc.env:HF_HOME}

config/backend/megatron_bridge/base.yaml

@@ -0,0 +1,7 @@
+# @package backend
+# Required per-experiment:
+megatron_path: ???
+hf_path: ???
+hf_model: ???
+tokenizer: ???
+extra_cli_args: []

config/backend/megatron_bridge/qwen3_0.6b.yaml

@@ -0,0 +1,4 @@
+# @package backend
+# Preset: Megatron checkpoint → HuggingFace Qwen3-0.6B with the same tokenizer.
+hf_model: "Qwen/Qwen3-0.6B"
+tokenizer: "Qwen/Qwen3-0.6B"

config/backend/megatron_bridge/qwen3_1.7b.yaml

@@ -0,0 +1,3 @@
+# @package backend
+hf_model: "Qwen/Qwen3-1.7B"
+tokenizer: "Qwen/Qwen3-1.7B"

config/backend/megatron_ftlauncher_localaddr.yaml

@@ -0,0 +1,50 @@
+# ft_launcher variant of the Megatron launcher (for fault-tolerant / warm-reserve
+# in-process restart). Bases on the `megatron` backend (empty torchrun_args) and
+# defines a clean ft_launcher arg set, rather than inheriting megatron_torchrun's
+# STATIC-rendezvous args (node_rank/master_addr/master_port) which ft_launcher
+# rejects and which can't simply be nulled (torchrun_args is typed dict[str,str]).
+#
+# Replaces `torchrun` with NVIDIA Resiliency Extension's `ft_launcher`
+# (https://github.com/NVIDIA/nvidia-resiliency-ext) — a torchrun-compatible
+# launcher that also starts a per-rank monitor (heartbeats + hang detection).
+# REQUIRED for Megatron's `--enable-ft-package` and its built-in
+# `FT_SIM_FAULT_DESC` fault simulator.
+#
+# In-process restart + warm reserve ("spare node jumps in") is driven by the
+# `--inprocess-*` Megatron args (set in the experiment); they use their own
+# TCPStore on MASTER_PORT+1/+2 and don't depend on the rendezvous backend here.
+#
+# VERIFIED ON JUPITER (job 569668, nvidia-resiliency-ext 0.5.0 in
+# MegatronTraining-JUPITER_pt2512_aarch64_202602231451.sif): launcher + c10d
+# rendezvous + reserve come up, the fault fires, the in-process restart runs
+# cleanly and the reserve node takes over and resumes training.
+#
+# DEPENDENCY: requires the 1-line Megatron fix that resets the FT rank-monitor
+# client between in-process restarts — add `ft_integration.shutdown()` to
+# `destroy_state()` in megatron/training/inprocess_restart.py. Without it every
+# restart hits AssertionError('rank monitor client is already initialized').
+defaults:
+  - megatron
+  - _self_
+
+# ft_launcher invocation. vs plain torchrun:
+#   * `torchrun` -> `ft_launcher`.
+#   * Rendezvous MUST be c10d (ft_launcher rejects rdzv_backend=static) -> use a
+#     shared rdzv_id and DON'T pass node_rank/master_addr/master_port.
+#   * FT section timeouts are hardcoded here (NOT routed through torchrun_args)
+#     because the cliargs resolver rewrites `_`->`-`. NVRX 0.5.0 names them
+#     `--ft-rank-section-timeouts` / `--ft-rank-out-of-section-timeout` (the
+#     `--ft-param-*` form from newer docs is rejected as unrecognized on 0.5.0).
+#     With `--calc-ft-timeouts` (experiment) Megatron refines + persists them.
+dist_cmd: "ft_launcher ${oc.join:'',${oc.mapcondtmpl:'\\-\\-.*','\\\n  %',' %',\
+ ${oc.cliargs:${.torchrun_args}}}}\\\n  --ft-rank-section-timeouts=setup:600,step:180,checkpointing:420 \\\n  --ft-rank-out-of-section-timeout=300 \\\n   ${.launcher_script}"
+
+torchrun_args:
+  nproc_per_node: ${slurm.sbatch.gpus_per_node}
+  nnodes: ${slurm.sbatch.nodes}
+  max_restarts: "3"     # job-level restart backstop; in-process handles per-fault recovery
+  tee: "3"
+  rdzv_backend: c10d    # ft_launcher requires c10d (static is rejected)
+  rdzv_endpoint: "$MASTER_ADDR:$MASTER_PORT"
+  rdzv_id: "$SLURM_JOB_ID"
+  local_addr: "$LOCAL_ADDR"

config/backend/oellm_eval.yaml

@@ -0,0 +1,29 @@
+defaults:
+  - oellm_eval: base
+  - _self_
+
+class_name: OELLMEvalBackend
+
+# CLI entrypoint (default assumes `oellm-eval` is on PATH from `uv tool install`).
+oellm_cmd: "oellm-eval"
+
+# oellm-autoexp owns the SLURM submission, so run oellm-evals's bash script
+# in-process via --local. Requires a venv with lm-eval/lighteval installed.
+local: true
+venv_path: ${oc.env:OELLM_EVAL_VENV,null}
+
+# Environment vars consumed by oellm-evals's generated eval script.
+# Path-dependent values (EVAL_BASE_DIR, EVAL_OUTPUT_DIR) intentionally omitted
+# here — set them in the experiment yaml so they pick up the right
+# job.base_output_dir for the eval stage (otherwise they resolve against the
+# experiment-level default during Hydra compose, not the sweep point).
+env:
+  HF_HOME: ${oc.env:HF_HOME}
+  HF_HUB_OFFLINE: "1"
+  QUEUE_LIMIT: "1"
+  GPUS_PER_NODE: "${slurm.sbatch.gpus_per_node}"
+  PARTITION: "${slurm.sbatch.partition}"
+  ACCOUNT: "${slurm.sbatch.account}"
+  SLURM_MEM: ${oc.env:SLURM_MEM,"96G"}
+  EVAL_CONTAINER_IMAGE: ""
+  SINGULARITY_ARGS: ""

config/backend/oellm_eval/base.yaml

@@ -0,0 +1,21 @@
+# @package backend
+
+# Default evaluation surface — override per experiment.
+models: []
+task_groups:
+  - open-sci-0.01
+tasks: null
+n_shot: null
+eval_csv_path: null
+
+# Runtime
+limit: null
+verbose: false
+download_only: false
+skip_checks: false
+trust_remote_code: true
+lm_eval_include_path: null
+slurm_template_var: null
+max_array_len: 128
+dry_run: false
+extra_cli_args: []

config/backend/oellm_eval/dclm.yaml

@@ -0,0 +1,3 @@
+# @package backend
+task_groups:
+  - dclm-core-22

config/backend/oellm_eval/opensci.yaml

@@ -0,0 +1,3 @@
+# @package backend
+task_groups:
+  - open-sci-0.01

config/container/jupiter_venv.yaml

@@ -0,0 +1,6 @@
+class_name: Venv
+venv_path: ${oc.env:OELLM_EVAL_VENV,${oc.env:HOME}/work/venv}
+
+env:
+  MACHINE_NAME: "JUPITER"
+  HF_HOME: ${oc.env:HF_HOME}

config/container/juwels_venv.yaml

@@ -0,0 +1,6 @@
+class_name: Venv
+venv_path: ${oc.env:OELLM_EVAL_VENV,${oc.env:HOME}/work/eval_venv}
+
+env:
+  MACHINE_NAME: "JUWELS"
+  HF_HOME: ${oc.env:HF_HOME}

config/container/leonardo_eval.yaml

@@ -0,0 +1,17 @@
+class_name: Container
+image: /leonardo_work/OELLM_prod2026/container_images/eval_env-leonardo.sif
+runtime: singularity
+# Includes /workspace/oellm-autoexp because the user-site oellm install
+# was an editable install made from inside the convert container, so its
+# `.pth` file points at /workspace/oellm-autoexp/submodules/oellm_evals/.
+bind:
+  - "/leonardo_scratch"
+  - "/leonardo"
+  - "/leonardo_work"
+  - "${oc.env:HOME}/work/Projects/oellm-autoexp:/workspace/oellm-autoexp"
+env:
+  WANDB_MODE: "offline"
+  MACHINE_NAME: "LEONARDO"
+  HF_HOME: ${oc.env:HF_HOME}
+  HF_HUB_OFFLINE: "1"
+python: "/usr/bin/python"

config/container/lumi_eval.yaml

@@ -0,0 +1,14 @@
+class_name: Container
+image: /scratch/project_462000963/containers/laif-rocm-6.4.4-pytorch-2.9.1-te-2.4.0-fa-2.8.0-triton-3.2.0.sif
+runtime: singularity
+bind:
+  - "/pfs/"
+  - "/scratch/"
+  - "/flash/"
+  - "${oc.env:HOME}"
+env:
+  WANDB_MODE: "offline"
+  MACHINE_NAME: "LUMI"
+  HF_HOME: ${oc.env:HF_HOME}
+  HF_HUB_OFFLINE: "1"
+python: "/usr/bin/python3"

config/experiments/korbi/chain_opensci_130M_train_eval_jupiter.yaml

@@ -0,0 +1,110 @@
+# @package _global_
+#
+# Chained multi-stage experiment: 1000-iter OpenSci 130M training → HF conversion
+# (in the train job) → standalone eval stage using OELLMEvalBackend.
+#
+# Stage layout:
+#   stage=train: Megatron training + same-job postprocess
+#                (torch_dist → torch → HF). Outputs HF checkpoint at
+#                ${job.base_output_dir}/hf/iter_0001000/.
+#   stage=eval:  OELLMEvalBackend in venv mode, waits via FileExistsCondition
+#                for the HF checkpoint, then runs oellm-evals on it.
+#
+# The two stages are submitted as separate SLURM jobs by oellm-autoexp;
+# the eval stage's start_condition blocks until the train stage produces
+# the expected file.
+#
+# We use OpenSci 130M because mcore_to_hf currently only ships an OpenSci
+# architecture profile. For Qwen3, add a Qwen3 profile under
+# oellm_autoexp/postprocess/resources/ and a new ckpt_convert_<arch>.yaml
+# postprocess pulling it in.
+
+defaults:
+  - ../base
+  - /backend: megatron_torchrun_localaddr
+  - /container: jupiter
+  - /slurm: jupiter
+  - /backend/megatron: [base_reference, opensci_ref_130M, data_mixturevitae]
+  - /postprocess: [ckpt_dist_to_torch_opensci, ckpt_convert_opensci]
+  - /sweep: none
+  - override /job: default
+  - _self_
+
+job:
+  name: "chain_opensci_130M_train_eval"
+  base_output_dir: ${oc.env:OUTPUT_DIR,"./output"}/${.name}_${index}
+
+# Default backend.megatron tweaks shared by the train stage.
+backend:
+  megatron:
+    legacy_tokenizer: false
+    train_iters: 1000
+    lr_warmup_iters: 50
+    lr_decay_iters: 1000
+    lr_wsd_decay_iters: 200
+    save_interval: 1000
+    eval_interval: 1000
+    ckpt_convert_format: torch
+    ckpt_convert_save: ${job.base_output_dir}
+    ckpt_step: ${backend.megatron.train_iters}
+    # Override data_mixturevitae default path (not present on this cluster).
+    # Re-use the small nemotron validation dataset for both train + valid;
+    # fine for a 1000-iter smoke test.
+    # split must be null when using train_data_path/valid_data_path (blend_per_split).
+    train_data_path:
+      - /e/scratch/projectnucleus/poeppel1/datasets/nemotron_downsampled_5p_valid/high-all
+    valid_data_path:
+      - /e/scratch/projectnucleus/poeppel1/datasets/nemotron_downsampled_5p_valid/high-all
+    data_path: []
+    split: null
+
+slurm:
+  env:
+    CUDA_DEVICE_MAX_CONNECTIONS: "1"
+  sbatch:
+    # The job-name validator expects slurm.name (with sweep index suffix);
+    # base config interpolates ${job.name} without it. Null lets the
+    # generator auto-emit the right name.
+    job_name: null
+    nodes: 1
+    time: "01:00:00"
+
+sweep:
+  type: list
+  filter: true
+  groups:
+    - type: list
+      configs:
+        - stage: train
+          # Train stage uses the dist→torch→HF postprocess inherited from
+          # defaults (run same-job after training). The eval stage clears
+          # postprocess via `postprocess: none` below.
+          # Override the default container image to one that's actually cached.
+          container.image: "${oc.env:CONTAINER_CACHE_DIR,.}/MegatronTraining-JUPITER_pt2512_aarch64_202602252238.sif"
+
+        - stage: eval
+          backend: oellm_eval           # swap backend config group
+          container: jupiter_venv        # swap container (venv, not apptainer)
+          slurm: jupiter_venv            # swap slurm (no apptainer wrapper)
+          postprocess: none              # swap to /postprocess/none.yaml → {}
+          # \\${...} delays interpolation until the DAG resolver injects the
+          # sibling.train.* context (same trick as test_diana.yaml). Lists with
+          # `${...}` items are emitted via Hydra's dotted-path pattern by
+          # param_to_cmdlines (key=[0], key.0=val) so the literal grammar
+          # doesn't choke on the `$`.
+          backend.models:
+            - "\\${sibling.train.job.base_output_dir}/hf/\\${oc.pad_iter:\\${sibling.train.backend.megatron.train_iters}}"
+          backend.task_groups:
+            - open-sci-0.01
+          backend.venv_path: "${oc.env:HOME}/work/venv"
+          backend.local: true
+          backend.skip_checks: true
+          backend.verbose: true
+          # Pin eval output to the eval stage's own base_output_dir.
+          backend.env.EVAL_BASE_DIR: "${oc.env:OUTPUT_DIR,./output}/${job.name}_1/eval"
+          backend.env.EVAL_OUTPUT_DIR: "${oc.env:OUTPUT_DIR,./output}/${job.name}_1/eval"
+          job.start_condition:
+            class_name: FileExistsCondition
+            path: "\\${sibling.train.job.base_output_dir}/hf/\\${oc.pad_iter:\\${sibling.train.backend.megatron.train_iters}}/model.safetensors"
+          slurm.sbatch.gpus_per_node: 1
+          slurm.sbatch.time: "00:30:00"