cwenzi's picture
|
download
raw
6.42 kB
metadata
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)

Xet Storage Details

Size:
6.42 kB
·
Xet hash:
10bd8819db608d77553a6b1d0f4deb83f2ca9d12b034f13c622bb35e6fdc449d

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.