# Lab Specification — Module FTDD-05: Axolotl

**Course**: Course 3 — LLM Fine-Tuning Masterclass
**Module**: FTDD-05 — Axolotl: Config-Driven Production Fine-Tuning
**Duration**: 35 minutes
**Environment**: Python 3.11+. A consumer GPU (RTX 4090 / 16–24GB) OR Apple Silicon (M-series) OR free Google Colab T4. ~6GB free disk. Multi-GPU is optional — this lab is designed to run single-GPU.

---

## Learning objectives

By the end of this lab you will have:

1. **Written an Axolotl YAML config from scratch** — all five sections (model, data, PEFT, hyperparams, distributed/output) — and run it to produce a trained model.
2. **Mutated the config** (LoRA on/off, batch size) and observed how a single declarative change propagates through the run — the config-as-source-of-truth pattern, felt directly.
3. **Run the reproducibility checklist** against your config and fixed any unspecified defaults or unpinned splits.
4. **Stated, in your own words, why the declarative YAML is more reproducible than a training script** — and what would force you off Axolotl onto raw TRL.

This lab is deliberately *config-first*. The point is to feel the recipe-as-file pattern before you spend the rest of the course deep in PEFT and alignment.

---

## Phase 0 — Environment setup (5 min)

```bash
python3.11 -m venv ftdd05-env && source ftdd05-env/bin/activate
# Axolotl pulls in TRL, transformers, peft, accelerate, deepspeed as dependencies.
pip install -q axolotl[flash-attn] transformers datasets torch
```

> **Note:** On Apple Silicon or Colab T4, omit `[flash-attn]` (FlashAttention is CUDA/Ampere+). A plain `pip install -q axolotl` works and falls back to standard attention. The lab is designed to run without FlashAttention.

Verify the install:

```bash
axolotl --version 2>/dev/null || python -c "import axolotl; print('axolotl OK')"
```

---

## Phase 1 — Write the YAML config (10 min)

Create `sft.yml`. This is the entire recipe — five sections. Read each comment.

```yaml
# sft.yml — Axolotl SFT on a 1B base
# ============================================

# --- 1. MODEL ---
base_model: openbmb/MiniCPM5-1B
tokenizer_type: AutoTokenizer

# --- 2. DATA ---
datasets:
  - path: HuggingFaceTB/smoltalk
    type: chat_template
    split: train
val_set_size: 0.02          # hold out 2% for eval (declared in the source of truth)

# --- 3. PEFT ---
adapter: qlora              # qlora | lora | none
lora_r: 8
lora_alpha: 16
lora_dropout: 0.05
lora_target_modules: [q_proj, v_proj]

# --- 4. TRAINING HYPERPARAMETERS (PIN THESE) ---
sequence_len: 2048
sample_packing: true
micro_batch_size: 1
gradient_accumulation_steps: 4    # effective batch = 1 * 4 = 4
num_epochs: 1
learning_rate: 2e-4               # PINNED — never rely on a default
lr_scheduler: cosine
warmup_steps: 10
bf16: auto                        # auto-detect; record what it resolves to

# --- 5. DISTRIBUTED + OUTPUT ---
# single-GPU: no deepspeed/fsdp block needed
output_dir: ./out/minicpm-smoltalk
save_steps: 100
eval_steps: 100
logging_steps: 10
```

Run it:

```bash
axolotl train sft.yml
```

**Record**: (a) the final training loss, (b) what `bf16: auto` resolved to (check the logs — it will state the detected precision), (c) the steerable-params percentage (Axolotl logs the LoRA param count).

> **What just happened (the teaching moment):** You wrote no training loop. The YAML IS the recipe. Axolotl read it, configured a TRL `SFTTrainer` with a QLoRA adapter, and ran. The same file would reproduce this run on a teammate's machine.

---

## Phase 2 — Mutate the config (10 min)

Now feel the declarative pattern: change ONE line, observe the propagation.

**Mutation A — LoRA off (full fine-tune):**

Edit `sft.yml`: change `adapter: qlora` to `adapter: none`. (You may also need to lower `micro_batch_size` if you OOM on full FT of 1B — but a 1B model usually fits.)

```bash
axolotl train sft.yml
```

**Record**: the new steerable-params percentage (should be ~100% now — all params train) and the final loss. Notice you changed ONE line; the entire training regime switched from QLoRA to full FT.

**Mutation B — change the batch recipe:**

Restore `adapter: qlora`. Now change `gradient_accumulation_steps: 4` to `gradient_accumulation_steps: 8` (and keep `micro_batch_size: 1`).

```bash
axolotl train sft.yml
```

**Record**: the effective batch size (now 1 × 8 = 8) and how the loss curve differs from Phase 1. Notice again: one line changed, the effective batch doubled, the rest of the recipe is identical.

> **The pattern:** every meaningful change is a localized edit to the YAML. There is no second script, no scattered code. The config is the complete source of truth.

---

## Phase 3 — The reproducibility checklist (5 min)

Run the six-item checklist against your `sft.yml`. For each, answer yes or fix it:

1. **Every hyperparameter pinned?** (Is any value riding an unspecified default? `learning_rate`, `lora_r`, `sequence_len` — all explicit?)
2. **Dataset path AND split explicit?** (`split: train` is declared — good. What about the val split?)
3. **Distributed strategy declared?** (Single-GPU: no block. That itself is a declared choice. If you were on multi-GPU, you'd need a `deepspeed:` or `fsdp:` block.)
4. **Precision known?** (What did `bf16: auto` resolve to? Record it. In production, pin it explicitly.)
5. **Eval declared?** (`val_set_size: 0.02`, `eval_steps: 100` — yes.)
6. **Reproducible from YAML alone?** (Could a teammate run `axolotl train sft.yml` and get your loss curve? Any hidden state?)

Fix anything that fails. Commit the final `sft.yml` to your report.

---

## Phase 4 — When would you leave Axolotl? (5 min)

No code. Answer in 3–5 sentences:

1. Your next job needs GRPO with a custom Python reward function (a domain-specific legal-citation scorer). Can you do this entirely in an Axolotl YAML? Why or why not? What do you reach for instead?
2. You have a single RTX 4090 and need to fine-tune a 7B model as fast as possible. Axolotl works, but what tool would give you ~2x throughput and ~half the VRAM, and why?
3. State in one sentence why "I tweaked the config live with CLI overrides" destroys the reproducibility guarantee, and what you must do instead.

---

## Deliverables

Submit `ftdd05-lab-report.md`:

- [ ] Phase 1: your `sft.yml` (all five sections); the final training loss; the detected precision; the LoRA steerable-params %.
- [ ] Phase 2: Mutation A (adapter: none) — new steerable-params % and loss; Mutation B (grad_accum 8) — effective batch size and loss-curve difference. Note that each was a one-line change.
- [ ] Phase 3: the six-item checklist with yes/fix for each; your final committed `sft.yml`.
- [ ] Phase 4: your 3–5 sentence answers.

---

## Solution key

- **Phase 1**: a valid `sft.yml` with all five sections. A successful run produces a decreasing loss curve. `bf16: auto` typically resolves to `true` on A100/H100/RTX 4090 and to FP16 on older GPUs / MPS. The LoRA steerable-params % should be **under 1%** (typically 0.1–0.5% for r=8 on q/v_proj of a 1B model) — the Steering Stack thesis (FT00) felt through Axolotl.
- **Phase 2**:
  - Mutation A (`adapter: none`): steerable params jumps to ~100% (all params train — full fine-tune). The loss curve may differ; on a small dataset full FT may overfit faster. The point is the ONE-LINE regime switch.
  - Mutation B (grad_accum 8): effective batch = 1 × 8 = 8 (doubled from Phase 1's 4). The loss curve should be smoother (larger batch reduces gradient noise) but the run takes ~2x longer per optimizer step. Again, a one-line change.
- **Phase 3**: model answers — (1) all hyperparams pinned (yes, if the student left no defaults); (2) split explicit (yes — `split: train`); (3) distributed declared (yes — single-GPU is a declared no-block choice); (4) precision known (the student should record what `bf16: auto` resolved to; in production, pin it); (5) eval declared (yes); (6) reproducible from YAML alone (yes, after fixes).
- **Phase 4** (model answers):
  1. No — a custom Python reward function is CODE, not config. Axolotl is configured, not programmed; a custom reward cannot be fully expressed in a static YAML. Reach for raw TRL's Python API (GRPOTrainer with the reward passed as a callable).
  2. Unsloth. Its hand-tuned Triton kernels replace TRL's attention/LoRA-backward kernels for ~2x throughput and ~half the VRAM on a single GPU, sub-30B. Axolotl works but leaves that speed on the table.
  3. Live overrides make the on-disk YAML no longer match what ran — the config becomes documentation, not source, and the run is unreproducible. You must write every override back into the YAML and commit it.

---

## Stretch goals

1. **Add a multi-GPU block.** If you have ≥2 GPUs, add a `deepspeed:` block pointing at a ZeRO-2 config (Axolotl ships examples in its repo) and re-run. Observe how the config declares sharding — this is the multi-GPU orchestration value-add over raw TRL. (Sets up the production path.)
2. **Swap to DPO via config.** Change the config to a DPO recipe over a small preference dataset (e.g., `trl-lib/ultrafeedback_binarized`) and run. Observe that the config surface is uniform across objectives — the same declarative pattern, a different trainer underneath. (Sets up FT13.)
3. **Diff two runs.** Run Phase 1, then change ONLY the learning rate (e.g., 2e-4 → 5e-5), commit both YAMLs to git, and `git diff` them. See reproducibility-as-a-diff — the exact property that makes config-as-source-of-truth win in production.