tinymind-native-colab-handoff/bundle/train/dpo_trainer.py

"""
DPO Trainer — Direct Preference Optimization

ทำให้โมเดลตอบดีขึ้นโดยการ align ด้วย preference pairs:
- chosen: คำตอบที่ดี (มี CoT reasoning)
- rejected: คำตอบที่แย่กว่า (ไม่มี reasoning / ผิด / สั้นเกิน)

อ้างอิง: "Direct Preference Optimization" (Rafailov et al. 2023)
"""

from __future__ import annotations

import json
import math
import time
from pathlib import Path
from functools import partial

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from tokenizers import Tokenizer

from model.config import OmegaConfig, small_config
from model.architecture import OmegaModel


# ─── DPO Dataset ──────────────────────────────────────────────────────────────

class DPODataset(Dataset):
    """โหลด DPO pairs จาก JSONL — แต่ละ record มี question, chosen, rejected"""

    CHAT_TEMPLATE = (
        "<bos><system>{system}</system>\n"
        "<user>{question}</user>\n"
        "<assistant>{response}<eos>"
    )
    DEFAULT_SYSTEM = (
        "คุณคือ TinyMind ผู้ช่วย AI ที่ฉลาด คิดอย่างเป็นระบบ และตอบถูกต้องเสมอ"
    )

    def __init__(self, path: str, tokenizer: Tokenizer, max_seq_len: int = 1024):
        self.tokenizer = tokenizer
        self.max_seq_len = max_seq_len
        self.records: list[dict] = []

        with open(path, encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                rec = json.loads(line)
                if rec.get("chosen") and rec.get("rejected") and rec.get("question"):
                    self.records.append(rec)

    def _encode(self, question: str, response: str, system: str) -> list[int]:
        text = self.CHAT_TEMPLATE.format(
            system=system, question=question, response=response
        )
        enc = self.tokenizer.encode(text)
        return enc.ids[: self.max_seq_len]

    def _prompt_len(self, question: str, system: str) -> int:
        prompt = (
            f"<bos><system>{system}</system>\n"
            f"<user>{question}</user>\n"
            f"<assistant>"
        )
        return len(self.tokenizer.encode(prompt).ids)

    def __len__(self) -> int:
        return len(self.records)

    def __getitem__(self, idx: int) -> dict:
        rec = self.records[idx]
        q = rec["question"]
        system = rec.get("system", self.DEFAULT_SYSTEM)

        chosen_ids  = self._encode(q, rec["chosen"],   system)
        rejected_ids = self._encode(q, rec["rejected"], system)
        prompt_len   = self._prompt_len(q, system)

        return {
            "chosen_ids":   chosen_ids,
            "rejected_ids": rejected_ids,
            "prompt_len":   prompt_len,
        }


def dpo_collate(batch: list[dict], pad_id: int = 0) -> dict:
    def pad(seqs: list[list[int]]) -> torch.Tensor:
        max_len = max(len(s) for s in seqs)
        return torch.tensor(
            [s + [pad_id] * (max_len - len(s)) for s in seqs], dtype=torch.long
        )

    return {
        "chosen_ids":   pad([b["chosen_ids"]   for b in batch]),
        "rejected_ids": pad([b["rejected_ids"]  for b in batch]),
        "prompt_len":   torch.tensor([b["prompt_len"] for b in batch], dtype=torch.long),
    }


# ─── DPO Loss ─────────────────────────────────────────────────────────────────

def compute_log_probs(
    model: OmegaModel,
    input_ids: torch.Tensor,
    prompt_len: torch.Tensor,
) -> torch.Tensor:
    """คำนวณ log prob ของส่วน response เท่านั้น (mask prompt out)"""
    out = model(input_ids)
    logits = out["logits"]  # [B, T, V]

    # shift: logits[t] predicts token[t+1]
    log_probs = F.log_softmax(logits[:, :-1, :], dim=-1)
    targets = input_ids[:, 1:]  # [B, T-1]

    # gather log prob ของ target tokens
    token_log_probs = log_probs.gather(2, targets.unsqueeze(-1)).squeeze(-1)  # [B, T-1]

    # mask: เก็บเฉพาะ response tokens (หลัง prompt)
    B, T = targets.shape
    mask = torch.zeros(B, T, device=input_ids.device, dtype=torch.bool)
    for i in range(B):
        start = min(int(prompt_len[i].item()), T)
        mask[i, start:] = True
    # ignore pad (id=0)
    mask &= targets != 0

    token_log_probs = token_log_probs * mask
    seq_log_probs = token_log_probs.sum(dim=-1) / mask.sum(dim=-1).clamp(min=1)
    return seq_log_probs  # [B]


def dpo_loss(
    policy_chosen_lp: torch.Tensor,
    policy_rejected_lp: torch.Tensor,
    ref_chosen_lp: torch.Tensor,
    ref_rejected_lp: torch.Tensor,
    beta: float = 0.1,
) -> torch.Tensor:
    """DPO loss: -log σ(β * (chosen_ratio - rejected_ratio))"""
    chosen_ratio   = policy_chosen_lp   - ref_chosen_lp
    rejected_ratio = policy_rejected_lp - ref_rejected_lp
    logits = beta * (chosen_ratio - rejected_ratio)
    loss = -F.logsigmoid(logits).mean()
    return loss


# ─── DPO Trainer ──────────────────────────────────────────────────────────────

DPO_CFG = {
    "data_path":       "data/filtered/dpo_pairs.jsonl",
    "tokenizer_path":  "data/tokenizer/tokenizer.json",
    "ref_checkpoint":  "checkpoints/omega_best.pt",   # frozen reference model
    "out_dir":         "checkpoints",
    "max_seq_len":     1024,
    "batch_size":      2,
    "grad_accum":      4,
    "lr":              5e-6,
    "beta":            0.1,
    "max_steps":       5_000,
    "save_every":      500,
    "dtype":           "bfloat16",
    "grad_clip":       1.0,
}


class DPOTrainer:
    def __init__(self, cfg: dict = DPO_CFG, model_cfg: OmegaConfig | None = None):
        self.cfg = cfg
        self.model_cfg = model_cfg or small_config()
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.dtype = torch.bfloat16 if cfg["dtype"] == "bfloat16" else torch.float16
        self.step = 0

    def setup(self):
        print(f"DPO Trainer | device={self.device} | dtype={self.dtype}")
        tok_path = self.cfg["tokenizer_path"]
        if not Path(tok_path).exists():
            raise FileNotFoundError(f"Tokenizer not found: {tok_path}")
        self.tokenizer = Tokenizer.from_file(tok_path)

        ds = DPODataset(self.cfg["data_path"], self.tokenizer, self.cfg["max_seq_len"])
        collate = partial(dpo_collate, pad_id=self.model_cfg.pad_token_id)
        self.loader = DataLoader(
            ds, batch_size=self.cfg["batch_size"],
            shuffle=True, collate_fn=collate, num_workers=2, pin_memory=True,
        )
        print(f"DPO dataset: {len(ds):,} pairs")

        # Policy model (trainable)
        ref_path = self.cfg.get("ref_checkpoint")
        if ref_path and Path(ref_path).exists():
            ckpt = torch.load(ref_path, map_location=self.device, weights_only=False)
            saved_cfg: OmegaConfig = ckpt["model_cfg"]
            self.policy = OmegaModel(saved_cfg).to(self.device)
            self.policy.load_state_dict(ckpt["model_state"])
            self.model_cfg = saved_cfg
        else:
            self.policy = OmegaModel(self.model_cfg).to(self.device)

        # Reference model (frozen copy)
        self.ref = OmegaModel(self.model_cfg).to(self.device)
        self.ref.load_state_dict(self.policy.state_dict())
        for p in self.ref.parameters():
            p.requires_grad_(False)
        self.ref.eval()

        self.policy.enable_grad_checkpointing()
        self.optimizer = torch.optim.AdamW(
            self.policy.parameters(), lr=self.cfg["lr"],
            betas=(0.9, 0.95), eps=1e-8, weight_decay=0.0
        )

    def train_step(self, batch: dict) -> float:
        self.policy.train()
        chosen_ids   = batch["chosen_ids"].to(self.device)
        rejected_ids = batch["rejected_ids"].to(self.device)
        prompt_len   = batch["prompt_len"].to(self.device)

        with torch.amp.autocast(device_type=self.device.type, dtype=self.dtype,
                                enabled=self.device.type == "cuda"):
            policy_chosen_lp   = compute_log_probs(self.policy, chosen_ids,   prompt_len)
            policy_rejected_lp = compute_log_probs(self.policy, rejected_ids, prompt_len)

            with torch.no_grad():
                ref_chosen_lp   = compute_log_probs(self.ref, chosen_ids,   prompt_len)
                ref_rejected_lp = compute_log_probs(self.ref, rejected_ids, prompt_len)

            loss = dpo_loss(
                policy_chosen_lp, policy_rejected_lp,
                ref_chosen_lp, ref_rejected_lp,
                beta=self.cfg["beta"],
            ) / self.cfg["grad_accum"]

        loss.backward()
        return loss.item() * self.cfg["grad_accum"]

    def save(self, tag: str = "dpo_latest"):
        path = Path(self.cfg["out_dir"]) / f"omega_{tag}.pt"
        torch.save({
            "step": self.step,
            "model_state": self.policy.state_dict(),
            "model_cfg": self.model_cfg,
        }, path)
        print(f"  Saved → {path}")

    def train(self):
        self.setup()
        data_iter = iter(self.loader)
        self.optimizer.zero_grad()
        t0 = time.time()
        running_loss = 0.0

        print(f"DPO training for {self.cfg['max_steps']:,} steps\n")
        while self.step < self.cfg["max_steps"]:
            for _ in range(self.cfg["grad_accum"]):
                try:
                    batch = next(data_iter)
                except StopIteration:
                    data_iter = iter(self.loader)
                    batch = next(data_iter)
                running_loss += self.train_step(batch)

            torch.nn.utils.clip_grad_norm_(self.policy.parameters(), self.cfg["grad_clip"])
            self.optimizer.step()
            self.optimizer.zero_grad()
            self.step += 1

            if self.step % 10 == 0:
                dt = time.time() - t0
                print(f"step {self.step:5d} | dpo_loss {running_loss/10:.4f} | {dt:.1f}s")
                running_loss = 0.0
                t0 = time.time()

            if self.step % self.cfg["save_every"] == 0:
                self.save(f"dpo_step{self.step}")

        self.save("dpo_final")
        print("DPO training complete!")


if __name__ == "__main__":
    trainer = DPOTrainer()
    trainer.train()