SKILL.md

---
name: llm-alignment-training
description: Implement post-training alignment (SFT, DPO, RLHF) for language models — from base model to conversational assistant. Covers masked instruction loss, preference optimization, reward modeling, and integration with custom training loops.
tags: [llm, alignment, sft, dpo, rlhf, fine-tuning, post-training]
---

# LLM Alignment Training

Implement SFT (Supervised Fine-Tuning), DPO (Direct Preference Optimization), and RLHF (Reinforcement Learning from Human Feedback) for language models. This skill covers the full pipeline from base pre-trained model to conversational, aligned assistant.

## When to Use

- User wants to add conversational ability to a base language model
- User wants to implement instruction-following training
- User wants to add preference-based alignment (DPO/RLHF)
- User has a custom model implementation and needs alignment training code
- User wants to understand masked loss strategies for SFT

## Core Concepts

### Training Stages

```
Pretrained Model → SFT → DPO/RLHF → Aligned Model
     (next-token)    (instruction)  (preference)
```

| Stage | Data | Loss Function | Purpose |
|:------|:-----|:--------------|:--------|
| SFT | (instruction, response) pairs | Masked cross-entropy | Learn to follow instructions |
| DPO | (instruction, chosen, rejected) triples | DPO loss | Learn to prefer better responses |
| RLHF | (instruction, chosen, rejected) + reward model | PPO/REINFORCE | Optimize for human preference |

### SFT: Masked Instruction Loss

The key insight: only compute loss on the **response** tokens, not the instruction tokens. This prevents the model from "re-learning" the question itself.

```
Input:  [bos] Question [ANS] Answer tokens...[eos]
Mask:   [0,0,0,0,0,0,0,0,0,1,1,1,1,1,1,1,1,1,1]
                                       ↑ loss starts here
```

Implementation:
1. Tokenize `instruction + response` as a single sequence
2. Build loss mask: 0 for instruction positions, 1 for response positions
3. Compute cross-entropy only where mask = 1
4. Normalize by number of response tokens

### DPO: Direct Preference Optimization

DPO eliminates the need for a separate reward model by using the policy model itself as the reward signal relative to a frozen reference model.

```
L_DPO = -log(σ(β × (log π(y_w)/π_ref(y_w) - log π(y_l)/π_ref(y_l))))
```

Where:
- `π` = policy model (being optimized)
- `π_ref` = reference model (frozen copy of SFT model)
- `y_w` = chosen (preferred) response
- `y_l` = rejected response
- `β` = temperature parameter (typically 0.1-0.5)

Key implementation details:
- Reference model is a frozen copy — never updated
- Both models share the same architecture
- Log probabilities are computed per-token and summed over the response
- The implicit reward is `r(y) = β × (log π(y) - log π_ref(y))`

### RLHF: Full Pipeline (when DPO is insufficient)

1. Train reward model on preference data
2. Optimize policy against reward model using PPO
3. KL penalty to prevent policy from diverging too far from reference

## Implementation Patterns

### For Custom C++ Implementations

When implementing alignment for a custom model (not PyTorch/HuggingFace):

1. **Reuse existing infrastructure**:
   - Model forward pass
   - Optimizer (AdamW)
   - Learning rate scheduler
   - Model serialization format

2. **Add alignment-specific components**:
   - Loss mask computation for SFT
   - Reference model copy for DPO
   - Log probability computation for DPO

3. **Data format**:
   - JSONL is the standard: `{"instruction": "...", "response": "..."}`
   - For DPO: `{"instruction": "...", "chosen": "...", "rejected": "..."}`
   - Alternative: plain text with delimiter (e.g., `|||`)

4. **Checkpoint format**:
   - Reuse existing model serialization
   - Save both policy and reference models for DPO

### Training Hyperparameters

| Parameter | SFT | DPO | RLHF |
|:----------|:----|:----|:-----|
| Learning rate | 1e-5 to 5e-5 | 1e-6 to 5e-6 | 1e-6 to 1e-5 |
| Epochs | 1-5 | 1-5 | 1-3 |
| Batch size | 4-32 | 4-32 | 4-16 |
| Max seq length | 512-4096 | 512-4096 | 512-2048 |
| Warmup ratio | 0.03-0.1 | 0.03-0.1 | 0.03-0.1 |
| β (DPO temp) | — | 0.1-0.5 | — |
| KL penalty | — | — | 0.01-0.1 |

## Common Pitfalls

### SFT Pitfalls

1. **Not masking instruction loss**: Model learns to "repeat" the question rather than answer it
2. **Too high learning rate**: Destroys pre-trained knowledge (catastrophic forgetting)
3. **Insufficient data diversity**: Model overfits to specific response patterns
4. **Wrong tokenization**: Instruction and response must use the same tokenizer as pre-training

### DPO Pitfalls

1. **Reference model not frozen**: Both models update, destroying the preference signal
2. **β too large**: Policy diverges too far from reference, quality degrades
3. **β too small**: No meaningful preference learning occurs
4. **Poor quality preference data**: Chosen/rejected pairs must have clear quality difference
5. **Not using SFT model as starting point**: DPO works best when initialized from SFT, not base model

### General Pitfalls

1. **Forgetting to call `model.train()` / `model.eval()`**: Batch norm and dropout behave differently
2. **Not saving checkpoints**: Alignment training can be unstable; save frequently
3. **No evaluation**: Always hold out eval data to monitor overfitting
4. **Ignoring sequence length**: Long sequences need more memory; consider gradient accumulation

## Verification Checklist

After implementing alignment training:

- [ ] SFT loss decreases over training
- [ ] DPO loss decreases (or accuracy of preference prediction increases)
- [ ] Model generates coherent responses to unseen instructions
- [ ] Model prefers chosen over rejected responses (DPO)
- [ ] Reference model remains unchanged (DPO)
- [ ] Checkpoints can be loaded and used for inference
- [ ] No NaN/Inf in loss or gradients
- [ ] Eval loss tracks training loss (no severe overfitting)

## Related Skills

- `huggingface-hub` — for using HuggingFace models/datasets
- `evaluating-llms-harness` — for benchmarking aligned models
- `serving-llms-vllm` — for serving aligned models
- `dspy` — for declarative LM program optimization

## References

- See `references/neuroflow-sft-dpo-implementation.md` for a concrete example of implementing SFT+DPO for a custom C++ LLM (NeuroFlow model with SN/ECN/DMN architecture)