Diagrams — Module FT11: The Training Loop with TRL

Module: FT11 — The Training Loop with TRL Diagram count: 5 Tool: Mermaid (primary). Each diagram validated in Mermaid Live Editor.


Diagram 1 — The SFTTrainer End-to-End Pipeline

Type: Linear pipeline Purpose: The single diagram that shows the full SFT job, in the order it runs on your machine. Six stages, from dataset to saved model. Reading the diagram: Left to right. Stages 1–3 are setup (data, model, PEFT). Stage 4 is config. Stage 5 is the training loop itself (the part that takes hours). Stage 6 is the artifact.

flowchart LR
  D["1. Dataset\nchat-format, train/eval split"] --> M["2. Model\nbf16 + FlashAttn2"]
  M --> P["3. PEFT config\nLoraConfig or None"]
  P --> C["4. SFTConfig\nTrainingArguments levers"]
  C --> T["5. trainer.train()\nloss curve · eval every N steps"]
  T --> S["6. Save / Merge / Load\nadapter or full model"]

  style D fill:#14141f,stroke:#5eead4,stroke-width:1.5px,color:#e4e4e8
  style M fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style P fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style C fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style T fill:#14141f,stroke:#5eead4,stroke-width:1.5px,color:#e4e4e8
  style S fill:#14141f,stroke:#5eead4,stroke-width:1.5px,color:#e4e4e8

Diagram 2 — The Loss-Curve Diagnostic Chart

Type: Comparison / diagnostic Purpose: The EKG of your training run. Three failure modes and the healthy baseline, plotted against steps. Reading the diagram: Each track is a loss-over-steps curve. Healthy descends and flattens. The three failures are the modes you must be able to name on sight. The annotation is the first thing to check.

flowchart TD
  subgraph Healthy["HEALTHY — descent + flatten"]
    H1["steep drop early\n→ gentle decline\n→ plateau (converged)"]
  end
  subgraph F1["FAILURE 1 — loss not decreasing"]
    N1["flat from step 1"]
    N2["check: LR too low?\ndata/template bug?"]
  end
  subgraph F2["FAILURE 2 — loss exploding (NaN)"]
    E1["spike to inf/nan"]
    E2["check: LR too high?\nFP16 overflow (use BF16)?\nEOS/template bug?"]
  end
  subgraph F3["FAILURE 3 — loss plateaus"]
    P1["descended, then flat"]
    P2["eval flat low = converged\neval rising = overfitting\n→ try LR warmup/decay, more data"]
  end

  Healthy -.->|"the target"| F1
  F1 -.->|"diagnose"| F2
  F2 -.->|"diagnose"| F3

  style Healthy fill:#14141f,stroke:#82e0aa,stroke-width:1.5px,color:#e4e4e8
  style H1 fill:#08080c,stroke:rgba(130,224,170,0.3),color:#82e0aa
  style F1 fill:#14141f,stroke:#f08080,color:#e4e4e8
  style N1 fill:#08080c,stroke:rgba(240,128,128,0.3),color:#f08080
  style N2 fill:#08080c,stroke:rgba(240,128,128,0.2),color:#9494a0
  style F2 fill:#14141f,stroke:#f08080,stroke-width:1.5px,color:#f08080
  style E1 fill:#08080c,stroke:rgba(240,128,128,0.3),color:#f08080
  style E2 fill:#08080c,stroke:rgba(240,128,128,0.2),color:#9494a0
  style F3 fill:#14141f,stroke:#f0a868,color:#e4e4e8
  style P1 fill:#08080c,stroke:rgba(240,168,104,0.3),color:#f0a868
  style P2 fill:#08080c,stroke:rgba(240,168,104,0.2),color:#9494a0

Diagram 3 — The TrainingArguments Decision Tree

Type: Decision tree Purpose: How to choose each lever. The config is not arbitrary — each knob answers a specific question about your hardware and method. Reading the diagram: Start at the top. Each branch is a decision you make when writing SFTConfig. The leaf is the setting and why.

flowchart TD
  Start["Writing SFTConfig"] --> Q1{PEFT or full FT?}
  Q1 -->|"LoRA/QLoRA (FT10)"| LR1["LR: 1e-4 to 5e-4\n(adapters learn faster)"]
  Q1 -->|"Full FT"| LR2["LR: 1e-5 to 5e-5\n(small steps, avoid forgetting)"]

  LR1 --> Q2{VRAM tight?}
  LR2 --> Q2
  Q2 -->|"Yes"| GC["gradient_checkpointing=True\n(~60-70% mem, ~30% slower)"]
  Q2 -->|"No"| GC2["gradient_checkpointing=False\n(full speed)"]

  GC --> Q3{batch fits?}
  GC2 --> Q3
  Q3 -->|"No"| GA["raise gradient_accumulation_steps\neffective_batch = bsz x accum x world"]
  Q3 -->|"Yes"| GA2["keep per_device_batch_size"]

  GA --> Q4{optimizer?}
  GA2 --> Q4
  Q4 -->|"QLoRA 4-bit base"| OPT["PagedAdamW (paged offload)"]
  Q4 -->|"memory-tight"| OPT2["AdamW 8-bit (bitsandbytes)"]
  Q4 -->|"default"| OPT3["AdamW"]

  Q5{hardware?} -->|"Ampere+"| BF["bf16=True (NOT fp16)"]
  Q5 -->|"pre-Ampere"| BF2["fp16 + grad clipping\n(plan to upgrade)"]

  style Start fill:#14141f,stroke:rgba(255,255,255,0.12),color:#e4e4e8
  style LR1 fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style LR2 fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style GC fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style GA fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style OPT fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style BF fill:#14141f,stroke:#5eead4,stroke-width:1.5px,color:#5eead4
  style BF2 fill:#14141f,stroke:#f08080,color:#f08080

Diagram 4 — The Save / Merge / Load Workflow

Type: Branching workflow Purpose: Three operations, three purposes. Confuse them and you ship the wrong artifact. Reading the diagram: After trainer.train(), you have a trained adapter (or full model). The branch you take depends on what you're doing next: experiment more, deploy a single model, or hot-swap at inference.

flowchart TD
  Train["trainer.train() done"] --> Save["trainer.save_model(path)"]
  Save --> Q{what next?}

  Q -->|"experiment: keep adapter separate"| Keep["adapter saved (few hundred MB)\nbase untouched · swappable"]
  Q -->|"deploy: ship one model"| Merge["merge_and_unload()\nadapter folded into base"]
  Q -->|"inference: hot-swap many adapters"| Hot["PeftModel.from_pretrained(base, adapter)\nload_adapter() to swap at runtime"]

  Merge --> QDeploy{serve how?}
  QDeploy -->|"vLLM / Ollama"| Serve["standalone model\n→ quantize (FT19) → serve (FT20)"]
  QDeploy -->|"Hub"| Push["merged.save_pretrained()\npush to Hub"]

  style Train fill:#14141f,stroke:#5eead4,stroke-width:1.5px,color:#e4e4e8
  style Save fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style Keep fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style Merge fill:#14141f,stroke:#5eead4,stroke-width:1.5px,color:#5eead4
  style Hot fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style Serve fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style Push fill:#14141f,stroke:#5eead4,color:#e4e4e8

Diagram 5 — The Eval-During-Training Loop

Type: Cyclic process Purpose: Why eval is not optional. The loop that catches overfitting and lets you pick the best checkpoint. Reading the diagram: The inner loop is training (forward, loss, backward, step). Every eval_steps, the loop pauses, runs eval on the held-out split, logs eval loss, and saves a checkpoint if it's the best so far. The gap between train and eval loss is the overfitting signal.

flowchart TD
  Train["TRAIN\nforward → loss → backward → optimizer step\nlog: train_loss, lr, grad_norm"] --> Q{every eval_steps?}
  Q -->|"no"| Train
  Q -->|"yes"| Eval["EVAL\nrun model on held-out split\nlog: eval_loss"]
  Eval --> Compare{eval_loss improved?}
  Compare -->|"yes"| Best["save checkpoint as 'best'\n(this is the one you ship)"]
  Compare -->|"no"| Check{train loss still dropping\neval rising?}
  Best --> Train
  Check -->|"yes: overfitting"| Stop["STOP at best checkpoint\nmemorizing, not steering"]
  Check -->|"no: converged or stuck"| Train

  style Train fill:#14141f,stroke:#5eead4,stroke-width:1.5px,color:#e4e4e8
  style Eval fill:#14141f,stroke:#5eead4,color:#5eead4
  style Best fill:#14141f,stroke:#82e0aa,color:#82e0aa
  style Stop fill:#14141f,stroke:#f08080,color:#f08080
  style Compare fill:#08080c,stroke:rgba(94,234,212,0.3),color:#e4e4e8
  style Check fill:#08080c,stroke:rgba(94,234,212,0.3),color:#e4e4e8
  style Q fill:#08080c,stroke:rgba(94,234,212,0.3),color:#e4e4e8

Validation notes

# Diagrams — Module FT11: The Training Loop with TRL

**Module**: FT11 — The Training Loop with TRL
**Diagram count**: 5
**Tool**: Mermaid (primary). Each diagram validated in [Mermaid Live Editor](https://mermaid.live).

---

## Diagram 1 — The SFTTrainer End-to-End Pipeline

**Type**: Linear pipeline
**Purpose**: The single diagram that shows the full SFT job, in the order it runs on your machine. Six stages, from dataset to saved model.
**Reading the diagram**: Left to right. Stages 1–3 are setup (data, model, PEFT). Stage 4 is config. Stage 5 is the training loop itself (the part that takes hours). Stage 6 is the artifact.

```mermaid
flowchart LR
  D["1. Dataset\nchat-format, train/eval split"] --> M["2. Model\nbf16 + FlashAttn2"]
  M --> P["3. PEFT config\nLoraConfig or None"]
  P --> C["4. SFTConfig\nTrainingArguments levers"]
  C --> T["5. trainer.train()\nloss curve · eval every N steps"]
  T --> S["6. Save / Merge / Load\nadapter or full model"]

  style D fill:#14141f,stroke:#5eead4,stroke-width:1.5px,color:#e4e4e8
  style M fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style P fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style C fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style T fill:#14141f,stroke:#5eead4,stroke-width:1.5px,color:#e4e4e8
  style S fill:#14141f,stroke:#5eead4,stroke-width:1.5px,color:#e4e4e8
```

---

## Diagram 2 — The Loss-Curve Diagnostic Chart

**Type**: Comparison / diagnostic
**Purpose**: The EKG of your training run. Three failure modes and the healthy baseline, plotted against steps.
**Reading the diagram**: Each track is a loss-over-steps curve. Healthy descends and flattens. The three failures are the modes you must be able to name on sight. The annotation is the first thing to check.

```mermaid
flowchart TD
  subgraph Healthy["HEALTHY — descent + flatten"]
    H1["steep drop early\n→ gentle decline\n→ plateau (converged)"]
  end
  subgraph F1["FAILURE 1 — loss not decreasing"]
    N1["flat from step 1"]
    N2["check: LR too low?\ndata/template bug?"]
  end
  subgraph F2["FAILURE 2 — loss exploding (NaN)"]
    E1["spike to inf/nan"]
    E2["check: LR too high?\nFP16 overflow (use BF16)?\nEOS/template bug?"]
  end
  subgraph F3["FAILURE 3 — loss plateaus"]
    P1["descended, then flat"]
    P2["eval flat low = converged\neval rising = overfitting\n→ try LR warmup/decay, more data"]
  end

  Healthy -.->|"the target"| F1
  F1 -.->|"diagnose"| F2
  F2 -.->|"diagnose"| F3

  style Healthy fill:#14141f,stroke:#82e0aa,stroke-width:1.5px,color:#e4e4e8
  style H1 fill:#08080c,stroke:rgba(130,224,170,0.3),color:#82e0aa
  style F1 fill:#14141f,stroke:#f08080,color:#e4e4e8
  style N1 fill:#08080c,stroke:rgba(240,128,128,0.3),color:#f08080
  style N2 fill:#08080c,stroke:rgba(240,128,128,0.2),color:#9494a0
  style F2 fill:#14141f,stroke:#f08080,stroke-width:1.5px,color:#f08080
  style E1 fill:#08080c,stroke:rgba(240,128,128,0.3),color:#f08080
  style E2 fill:#08080c,stroke:rgba(240,128,128,0.2),color:#9494a0
  style F3 fill:#14141f,stroke:#f0a868,color:#e4e4e8
  style P1 fill:#08080c,stroke:rgba(240,168,104,0.3),color:#f0a868
  style P2 fill:#08080c,stroke:rgba(240,168,104,0.2),color:#9494a0
```

---

## Diagram 3 — The TrainingArguments Decision Tree

**Type**: Decision tree
**Purpose**: How to choose each lever. The config is not arbitrary — each knob answers a specific question about your hardware and method.
**Reading the diagram**: Start at the top. Each branch is a decision you make when writing `SFTConfig`. The leaf is the setting and why.

```mermaid
flowchart TD
  Start["Writing SFTConfig"] --> Q1{PEFT or full FT?}
  Q1 -->|"LoRA/QLoRA (FT10)"| LR1["LR: 1e-4 to 5e-4\n(adapters learn faster)"]
  Q1 -->|"Full FT"| LR2["LR: 1e-5 to 5e-5\n(small steps, avoid forgetting)"]

  LR1 --> Q2{VRAM tight?}
  LR2 --> Q2
  Q2 -->|"Yes"| GC["gradient_checkpointing=True\n(~60-70% mem, ~30% slower)"]
  Q2 -->|"No"| GC2["gradient_checkpointing=False\n(full speed)"]

  GC --> Q3{batch fits?}
  GC2 --> Q3
  Q3 -->|"No"| GA["raise gradient_accumulation_steps\neffective_batch = bsz x accum x world"]
  Q3 -->|"Yes"| GA2["keep per_device_batch_size"]

  GA --> Q4{optimizer?}
  GA2 --> Q4
  Q4 -->|"QLoRA 4-bit base"| OPT["PagedAdamW (paged offload)"]
  Q4 -->|"memory-tight"| OPT2["AdamW 8-bit (bitsandbytes)"]
  Q4 -->|"default"| OPT3["AdamW"]

  Q5{hardware?} -->|"Ampere+"| BF["bf16=True (NOT fp16)"]
  Q5 -->|"pre-Ampere"| BF2["fp16 + grad clipping\n(plan to upgrade)"]

  style Start fill:#14141f,stroke:rgba(255,255,255,0.12),color:#e4e4e8
  style LR1 fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style LR2 fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style GC fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style GA fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style OPT fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style BF fill:#14141f,stroke:#5eead4,stroke-width:1.5px,color:#5eead4
  style BF2 fill:#14141f,stroke:#f08080,color:#f08080
```

---

## Diagram 4 — The Save / Merge / Load Workflow

**Type**: Branching workflow
**Purpose**: Three operations, three purposes. Confuse them and you ship the wrong artifact.
**Reading the diagram**: After `trainer.train()`, you have a trained adapter (or full model). The branch you take depends on what you're doing next: experiment more, deploy a single model, or hot-swap at inference.

```mermaid
flowchart TD
  Train["trainer.train() done"] --> Save["trainer.save_model(path)"]
  Save --> Q{what next?}

  Q -->|"experiment: keep adapter separate"| Keep["adapter saved (few hundred MB)\nbase untouched · swappable"]
  Q -->|"deploy: ship one model"| Merge["merge_and_unload()\nadapter folded into base"]
  Q -->|"inference: hot-swap many adapters"| Hot["PeftModel.from_pretrained(base, adapter)\nload_adapter() to swap at runtime"]

  Merge --> QDeploy{serve how?}
  QDeploy -->|"vLLM / Ollama"| Serve["standalone model\n→ quantize (FT19) → serve (FT20)"]
  QDeploy -->|"Hub"| Push["merged.save_pretrained()\npush to Hub"]

  style Train fill:#14141f,stroke:#5eead4,stroke-width:1.5px,color:#e4e4e8
  style Save fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style Keep fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style Merge fill:#14141f,stroke:#5eead4,stroke-width:1.5px,color:#5eead4
  style Hot fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style Serve fill:#14141f,stroke:#5eead4,color:#e4e4e8
  style Push fill:#14141f,stroke:#5eead4,color:#e4e4e8
```

---

## Diagram 5 — The Eval-During-Training Loop

**Type**: Cyclic process
**Purpose**: Why eval is not optional. The loop that catches overfitting and lets you pick the best checkpoint.
**Reading the diagram**: The inner loop is training (forward, loss, backward, step). Every `eval_steps`, the loop pauses, runs eval on the held-out split, logs eval loss, and saves a checkpoint if it's the best so far. The gap between train and eval loss is the overfitting signal.

```mermaid
flowchart TD
  Train["TRAIN\nforward → loss → backward → optimizer step\nlog: train_loss, lr, grad_norm"] --> Q{every eval_steps?}
  Q -->|"no"| Train
  Q -->|"yes"| Eval["EVAL\nrun model on held-out split\nlog: eval_loss"]
  Eval --> Compare{eval_loss improved?}
  Compare -->|"yes"| Best["save checkpoint as 'best'\n(this is the one you ship)"]
  Compare -->|"no"| Check{train loss still dropping\neval rising?}
  Best --> Train
  Check -->|"yes: overfitting"| Stop["STOP at best checkpoint\nmemorizing, not steering"]
  Check -->|"no: converged or stuck"| Train

  style Train fill:#14141f,stroke:#5eead4,stroke-width:1.5px,color:#e4e4e8
  style Eval fill:#14141f,stroke:#5eead4,color:#5eead4
  style Best fill:#14141f,stroke:#82e0aa,color:#82e0aa
  style Stop fill:#14141f,stroke:#f08080,color:#f08080
  style Compare fill:#08080c,stroke:rgba(94,234,212,0.3),color:#e4e4e8
  style Check fill:#08080c,stroke:rgba(94,234,212,0.3),color:#e4e4e8
  style Q fill:#08080c,stroke:rgba(94,234,212,0.3),color:#e4e4e8
```

---

## Validation notes

- All five diagrams use the course design system colors: `#14141f` panel fill, `#5eead4` accent for primary, `#82e0aa` (ok) / `#f08080` (danger) / `#f0a868` (warn) for semantic emphasis, `rgba(255,255,255,0.12)` for secondary borders, `#e4e4e8` / `#9494a0` for text.
- Paste each into [Mermaid Live Editor](https://mermaid.live) to render. All use stable Mermaid syntax (`flowchart` with `TD`/`LR`, `subgraph`) supported in current Mermaid (v10.4+).
- For the slide deck (artifact 03), these are rendered as static captures from Mermaid Live, inlined into reveal.js.