Skip to content

alignment.md

Latest commit

 

History

History
244 lines (174 loc) · 9.49 KB

alignment.md

File metadata and controls

244 lines (174 loc) · 9.49 KB
Raw

Alignment & Post-Training Guide

ForgeLM supports the complete modern post-training stack: SFT → Preference Optimization → Reasoning RL. This guide explains when and how to use each method.

Method Overview

Method trainer_type Dataset Format When to Use
SFT "sft" System/User/Assistant or messages Instruction tuning — teach the model what to say
DPO "dpo" chosen / rejected pairs Preference alignment — teach how to say it better
SimPO "simpo" chosen / rejected pairs Like DPO but no reference model (lower memory)
KTO "kto" completion + label (bool) Binary feedback — only thumbs up/down available
ORPO "orpo" chosen / rejected pairs SFT + alignment in one stage
GRPO "grpo" prompt only Reasoning RL — model generates and self-improves

The Modern Post-Training Stack

Most production LLMs in 2026 follow this pipeline:

Base Model
    ↓
[Stage 1] SFT — instruction tuning on curated data
    ↓
[Stage 2] DPO/SimPO/KTO — preference alignment
    ↓
[Stage 3] GRPO — reasoning RL (optional, for math/code)
    ↓
Production Model

ForgeLM handles each stage as a separate forgelm run with different configs.

Stage 1: Supervised Fine-Tuning (SFT)

Goal: Teach the model to follow instructions in your domain.

Dataset Format

{"System": "You are a legal assistant.", "User": "What is a tort?", "Assistant": "A tort is a civil wrong..."}

Or the modern messages format:

{"messages": [{"role": "system", "content": "..."}, {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}

Config

model:
  name_or_path: "meta-llama/Llama-3.1-8B-Instruct"
  load_in_4bit: true

lora:
  r: 16
  alpha: 32
  target_modules: ["q_proj", "k_proj", "v_proj", "o_proj"]

training:
  trainer_type: "sft"
  num_train_epochs: 3
  learning_rate: 2.0e-5
  per_device_train_batch_size: 4

data:
  dataset_name_or_path: "./data/sft_data.jsonl"
forgelm --config sft_config.yaml

Stage 2: Preference Alignment

After SFT, align the model's responses with human preferences. Choose based on your data:

DPO — Direct Preference Optimization

Best for: You have paired preference data (chosen vs rejected responses).

{"prompt": "Explain recursion", "chosen": "Recursion is a technique where...", "rejected": "Recursion means doing something again..."}
training:
  trainer_type: "dpo"
  dpo_beta: 0.1          # Temperature — lower = stronger preference signal
  learning_rate: 5.0e-6   # Lower LR than SFT
  num_train_epochs: 1     # 1-2 epochs usually sufficient

data:
  dataset_name_or_path: "./data/preferences.jsonl"

SimPO — Simple Preference Optimization

Best for: Same data as DPO, but you want lower memory (no reference model needed).

SimPO outperforms DPO at 7B+ scale (+6.4 points on AlpacaEval 2).

training:
  trainer_type: "simpo"
  simpo_beta: 2.0         # Scaling parameter
  simpo_gamma: 0.5        # Margin term
  learning_rate: 5.0e-6

KTO — Kahneman-Tversky Optimization

Best for: You only have binary feedback (thumbs up/down), not paired preferences. More practical for production data collection.

{"prompt": "What is Python?", "completion": "Python is a programming language...", "label": true}
{"prompt": "What is Python?", "completion": "Python is a snake.", "label": false}
training:
  trainer_type: "kto"
  kto_beta: 0.1
  learning_rate: 5.0e-6

data:
  dataset_name_or_path: "./data/kto_feedback.jsonl"

ORPO — Single-Stage SFT + Alignment

Best for: You want to combine SFT and alignment in one training run. Uses chosen/rejected data but also learns from the instruction format.

training:
  trainer_type: "orpo"
  orpo_beta: 0.1

Stage 3: Reasoning RL (GRPO)

Best for: Math, code, reasoning tasks where outputs can be verified. This is the method behind DeepSeek-R1.

GRPO generates multiple responses per prompt, scores them, and reinforces better ones — no human preference data needed.

{"prompt": "Solve: What is 15% of 240?", "gold_answer": "36"}
training:
  trainer_type: "grpo"
  grpo_num_generations: 4              # Generate 4 responses per prompt
  grpo_max_completion_length: 512      # Max tokens per completion (legacy alias `grpo_max_new_tokens` still accepted)
  grpo_reward_model: null    # See "Reward selection" below.
  learning_rate: 1.0e-6      # Very low LR for RL stability
  num_train_epochs: 1

data:
  dataset_name_or_path: "./data/math_prompts.jsonl"

Reward selection

GRPO needs a reward signal. ForgeLM wires reward callables additively (TRL sums multiple reward funcs into a single scalar):

  1. grpo_reward_model set — Loads the HF sequence-classification model at that path and uses its scalar output as the only reward signal. The built-in rewards below are bypassed; the operator opted into a learned reward.
  2. No grpo_reward_model — A baseline reward is always wired:
    • combined_format_length_reward (forgelm/grpo_rewards.py) — 0.8 × format_match + 0.2 × length_shaping. The format component returns 1.0 when the generation ends with Answer: <value> (case-insensitive, units allowed); the length component returns min(len(completion) / 200, 1.0) so early training has a non-flat gradient even before format compliance kicks in.
    • _math_reward_fn (forgelm/trainer.py) — appended only when the dataset has a gold_answer field. Captures the value after the last Answer: marker (so a self-correcting generation is graded on the answer it actually concludes with, consistent with the end-anchored format reward), strips common units ($, %, km/h, m², liters, …), and compares to gold_answer with exact-string match first, then numeric tolerance (1e-6). Returns 1.0 for a correct answer, 0.0 otherwise.

The bundled forgelm quickstart grpo-math template ships with gold_answer populated, so the model gets both format teaching AND correctness teaching out of the box. To use a real reward model on top of grpo-math, set grpo_reward_model and the built-in rewards are bypassed.

For your own dataset: the format+length baseline applies regardless. Add a gold_answer field per row to also get the correctness signal — the prompt's expected output format is Answer: <value> (with optional units that get stripped).

Note: GRPO requires a reward function or verifiable reward. For math, correctness of the answer is the reward. For general text, you may need a reward model.

Reward-model memory footprint. When grpo_reward_model is set, that model is loaded alongside the policy model for the whole run, so plan VRAM for both. ForgeLM loads it at the same compute dtype the rest of the pipeline resolves — bf16 where supported, else fp16 — rather than the checkpoint default of fp32, because a full-precision reward model sitting beside a 4-bit policy model is a realistic single-GPU OOM path. On a CPU-only host it falls back to fp32, since PyTorch does not implement fp16 matmul on CPU. See forgelm/trainer.py (the AutoModelForSequenceClassification.from_pretrained call in the GRPO reward path).

There is no 4-bit option for the reward model, deliberately, and none is planned. model.load_in_4bit applies to the policy model only. Quantising the reward model perturbs the very scalar GRPO optimises against — a shifted reward scale changes the objective rather than just the memory profile, and the resulting run is not comparable to an unquantised one. If reward-model VRAM is the binding constraint, use a smaller reward checkpoint or move to the built-in rewards above; do not expect a quantisation knob.

Choosing the Right Method

Do you have paired preferences (chosen/rejected)?
├── Yes → Is memory a concern?
│   ├── Yes → SimPO
│   └── No → DPO
├── No → Do you have binary feedback (good/bad)?
│   ├── Yes → KTO
│   └── No → Do you have verifiable rewards (math/code)?
│       ├── Yes → GRPO
│       └── No → Just use SFT

ForgeLM's --wizard mode helps you choose:

forgelm --wizard
# Step 4 asks: "Choose your training objective"
# Shows format requirements for each method

Multi-Stage Pipeline Example

# Stage 1: SFT
forgelm --config configs/stage1_sft.yaml

# Stage 2: DPO (uses the SFT model as base)
# In stage2_dpo.yaml, set:
#   model.name_or_path: "./checkpoints_sft/final_model"
forgelm --config configs/stage2_dpo.yaml

# Stage 3: GRPO (uses the DPO model as base)
forgelm --config configs/stage3_grpo.yaml

Available since v0.7.0 (Phase 14): The pipeline: config block chains multi-stage training (SFT → DPO → GRPO) in a single YAML. See Multi-Stage Training Pipelines for the full guide.

Tips

  • Learning rate: SFT uses 1e-5 to 3e-5. Alignment methods use 5e-7 to 5e-6. GRPO uses 1e-6 or lower.
  • Epochs: SFT typically needs 2-3 epochs. Alignment methods usually need 1-2 epochs. More is not better.
  • Data quality > data quantity: 1,000 high-quality preference pairs often outperform 50,000 noisy ones.
  • Always evaluate: Use auto_revert: true with max_acceptable_loss to catch quality regressions.
  • Scale matters: Research (arxiv 2603.19335) shows algorithm rankings are scale-dependent — SimPO is best at 7B but DPO may be better at 1.5B.