axolotl-ota-sft-integration

# Axolotl ⇄ OpenThoughts-Agent SFT-backend integration — overview **Status:** COMPLETE + merged to `penfever/working` (merge `02d676d0`, 2026-07-02). **Where it ran:** TACC Vista (GH200, aarch64), conda env `sft-axolotl`. **One-line result:** OpenThoughts-Agent can now run SFT through **axolotl** (`--sft_backend axolotl`) as a drop-in alternative to LLaMA-Factory, with the **delphi** chat-template masking **validated** on the real delphi path (jinja-as-ground-truth: train == serve). --- ## 1. What was done Two staged design→execute cycles (plans + per-stage scopes under `notes/`): **Cycle 1 — marin axolotl fork with 3 ported features** (`notes/cycle1_marin-fork-3feature-port/`). Created `marin-community/axolotl @ feuer/marin-fork-3feature-port` (`3c206072`, axolotl 0.17.0.dev0): - `delphi.jinja` chat template (auto-globbed) — the Llama-3 turn format + reasoning/tool tokens (`<|start_think|>`/`<|end_think|>`/`<|tool_call|>`/`<|tool_result|>`). - `template_integrity` plugin — the **save-time footgun fix**: keeps a populated `chat_template` embedded in `tokenizer_config.json` on every checkpoint save (`tokenizer_save_jinja_files: false` + per-checkpoint coverage) so the model is not silently OOD at serve time (the historical 0%-SWE-bench bug). - `mfu` plugin (MFU logging) + `supabase_registry` plugin (opt-in train-end model registration). **Cycle 2 — wire axolotl into `hpc.launch` as an SFT backend** (`notes/cycle2_sft-launch-backend/`). `--sft_backend {llamafactory,axolotl}` (LF default, flag-off byte-identical). Submodule `sft/axolotl` pinned at `3c206072`. Runner dispatch (`-m axolotl.cli.train`), an LF-exp-args → axolotl-YAML translator (`hpc/axolotl_config_utils.py`), config assets (`sft/axolotl_configs/`), and validation gates (`sft/axolotl_gates/`). Validated on TACC across 6 stages. --- ## 2. Gate results (see `results/` for the raw artifacts) | Gate | Verdict | Evidence | |---|---|---| | Stage 3 — end-to-end smoke | ✅ GO | job 801458; loss 2.272→0.951 (6 steps); SDPA; ckpt w/ embedded delphi template | | Stage 4 — footgun-through-launcher | ✅ GO | embedded `chat_template` **byte-identical** to canonical `delphi.jinja` across output + checkpoint dirs | | Stage 6 — LF-vs-axolotl loss-match | ❌ NO-GO (understood) | axolotl 1.567→**1.549**, LF 1.334→**0.702** (both 30 steps); 120% gap = a **framework turn-masking divergence on the non-canonical guanaco/llama3 dataset**, not a defect | | **Delphi masking canary** (the check that matters) | ✅ **PASS** | job 802053; loss 2.628→2.532 (20 steps); `<\|start_think\|>…<\|end_think\|>`+assistant **trained**, user/system **masked**, **0** "Last turn is not trainable" skips; trainable fraction 73–96% (via `axolotl.cli.preprocess --debug`) | **Takeaway:** the Stage-6 guanaco gap was a masking divergence *specific to that non-canonical dataset*. On the **real delphi path** the axolotl `delphi` template masks correctly — jinja-as-ground-truth holds. ### Delphi mask-dump evidence (job 802053, `axolotl.cli.preprocess --debug`) Dump format `text(label, token_id)`: label `-100` = masked, label == token_id = trained. Consistent across 4/4 examples: | Assertion | Result | |---|---| | `<\|start_think\|>`/`<\|end_think\|>` reasoning span trained | `(128002,128002)` + `(128003,128003)` **TRAINED** | | assistant answer trained | `The(791,791) answer(4320,4320) is(374,374)…` **TRAINED** | | user turns masked | `user(-100,882)` + content `(-100,…)` **MASKED** | | system turns masked | `<\|begin_of_text\|>`, system header+content `(-100,…)` **MASKED** | | assistant header masked (only content trained) | `assistant(-100,78191)` **MASKED** | | Llama-3 structure segmented | `<\|start_header_id\|>…<\|eot_id\|>` boundaries correct | | `Last turn is not trainable, skipping` | **0 occurrences** | Trainable-token fractions per example: **659/763 (86%), 1420/1474 (96%), 636/726 (88%), 463/631 (73%)** — high, as expected for reasoning data (short masked prompt, long trained reasoning+answer). One benign nuance: the trailing `<|eot_id|>` of the last assistant turn is masked (a `train_on_eos` policy choice, not a defect). Loss series (from `results/trainer_states/`): - `axolotl_parity_llama3_ckpt30.json` — 30 steps, 1.5667→1.5490 - `lf_parity_llama3_ckpt30.json` — 30 steps, 1.3342→0.7020 - `delphi_canary_ckpt20.json` — 20 steps, 2.6280→2.5315 - `axolotl_smoke_ckpt6.json` — 6 steps, 2.2722→0.9505 --- ## 3. Fixes shipped along the way **Axolotl backend / launcher (7):** submodule + `--sft_backend` selector; runner dispatch; LF→axolotl config translator; CLI `global_batch_size` int-cast; **short job-scoped TMPDIR** (AF_UNIX 108-byte `sun_path` fix, length-conditional, BOTH backends); **`expandable_segments`** allocator (BOTH backends); the `torchao==0.17.0` aarch64 env fix (doc, not a submodule patch). **LLaMA-Factory transformers-5.x (4, bonus — unbreak LF SFT on modern transformers):** LF-fork `d20b8666` = `add_special_tokens(replace_additional_special_tokens=…)` signature-guard; launcher: `report_to=none` on no-internet nodes (wandb-0.28 service-socket crash), the TMPDIR + expandable_segments generalizations above. LF pin bumped `6617a420 → d20b8666`. --- ## 4. Example launch commands ### 4a. Axolotl SFT — the validated **delphi masking canary** (TACC Vista, aarch64) ```bash # prereq once: prep the Delphi tokenizer (single delphi tokens) python sft/delphi/prepare_delphi_tokenizer.py \ --model laion/delphi-3e18-p33m67-k0p20-lr83-a003 --output $SCRATCH/delphi-canary-tok # launch (WANDB_MODE=disabled: TACC is an internet node -> report_to=wandb; wandb 0.28 crashes offline) WANDB_MODE=disabled python -m hpc.launch --job_type sft --sft_backend axolotl \ --train_config_path sft/axolotl_configs/delphi_canary.yaml \ --model_path $SCRATCH/delphi-canary-tok --conda_env sft-axolotl \ --dataset laion/llama-nemotron-science-reasoning-on-le3000tok-100k-canonical-think \ --messages messages --role_tag role --content_tag content \ --partition gh-dev --num_nodes 1 --gpus_per_node 1 --time_limit 01:00:00 ``` > **Launcher gotcha:** `hpc.launch` REBUILDS the axolotl `datasets:` block from `--dataset` + > `--messages/--role_tag/--content_tag`; the in-config `datasets:` block is honored ONLY by direct > `axolotl.cli.preprocess`. So pass those dataset flags on the launcher path. ### 4b. Mask-dump validation (the decisive masking check — no training) ```bash python -m axolotl.cli.preprocess sft/axolotl_configs/delphi_canary.yaml \ --base_model $SCRATCH/delphi-canary-tok --debug # via srun (login node OOMs on tokenizer loads) # inspect: <|start_think|>..<|end_think|> + assistant TRAINED (label==id), user/system MASKED (-100), # 0 "Last turn is not trainable, skipping" ``` ### 4c. Axolotl vs LLaMA-Factory SFT (backend swap — same launcher, same flags) ```bash # axolotl backend python -m hpc.launch --job_type sft --sft_backend axolotl --conda_env sft-axolotl \ --train_config_path sft/axolotl_configs/<cfg>.yaml --dataset <ds> --messages messages \ --role_tag role --content_tag content --num_nodes 1 --gpus_per_node 1 --time_limit 02:00:00 # llamafactory backend (default; omit --sft_backend) python -m hpc.launch --job_type sft --sft_backend llamafactory --conda_env otagent \ --train_config_path sft/lf_configs/<family>/<cfg>.yaml --dataset <ds> \ --role_tag role --user_tag user --assistant_tag assistant --content_tag content \ --num_nodes 1 --gpus_per_node 1 --time_limit 02:00:00 ``` ### 4d. Downstream — RL the SFT'd delphi model on CoreWeave / Iris (the next pipeline stage) The axolotl SFT backend produces the instruction-tuned delphi checkpoint that then goes to GRPO RL on the CoreWeave H100 GPU cluster (via the marin Iris SDK). Representative (see skill `rl-agentic-launch-iris`): ```bash python -m rl.cloud.launch_rl_iris \ --rl_config <rl_yaml> --model_path laion/<sft-delphi-ckpt> \ --train_data <task_parquet> --num-nodes 1 \ --rendezvous-dir <gs://…/rendezvous> --job-name <name> --priority normal --cpu 48 --max-retries 1 ``` > This SFT→RL handoff is why the delphi masking correctness matters: train==serve on `delphi.jinja` means > the SFT'd model isn't OOD when it hits the RL rollouts + eval. --- ## 5. Reproduce / investigate — directory guide - `OVERVIEW.md` — this file. - `agent_logs/` — the full dated execute logs for both cycles (blow-by-blow debug record). - `notes/cycle1_*`, `notes/cycle2_*` — the staged design plans + per-stage scopes + the DEFERRED-TACC checklist. - `configs/` — the axolotl configs (`delphi_canary.yaml`, `smoke.yaml`, `axolotl_parity_llama3.yaml`, `marin_delphi_all3.yaml`), the LF parity config, and the two validation gate scripts (`stage4_footgun_through_launcher.py`, `stage6_parity.py`). - `results/trainer_states/` — the loss series (JSON) for smoke / both parity sides / the delphi canary. - `results/rendered_configs/` — the exact launcher-rendered axolotl train config for the canary (byte repro). - `results/run_logs/` — the training `.out` logs (env fixes engaged, SDPA, template embed, loss lines) for the axolotl parity (801488), LF parity (801998), and the delphi canary (802053). ## 6. Key references (in-repo) - Backend wiring: `hpc/sft_launch_utils.py`, `hpc/axolotl_config_utils.py`, `hpc/arguments.py`. - Configs/gates: `sft/axolotl_configs/`, `sft/axolotl_gates/`, `sft/delphi/prepare_delphi_tokenizer.py`, `sft/delphi/dataset_info.json`. - Dependency facts + gotchas: `.claude/projects/axolotl/axolotl.md`. - Skill: `.claude/skills/sft-launch/` (merged jupiter+leonardo; backend + delphi guidance). - Fork: `marin-community/axolotl @ feuer/marin-fork-3feature-port` (`3c206072`).