code/train_dpo.py

# ==============================================================================
# 🚀 ViuAI Sarus-500M — Direct Preference Optimization (DPO) Training Engine
# ==============================================================================
# Ultra-optimized Native PyTorch DPO Engine with Multi-GPU DDP & Single-GPU Support
# Uses the SFT checkpoint selected by --sft_version as base policy + frozen reference.
# ==============================================================================

import os
import sys
import json
import math
import time
import shutil
import random
import argparse
from typing import Dict, List, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
if hasattr(sys.stderr, "reconfigure"):
    sys.stderr.reconfigure(encoding="utf-8", errors="replace")

# ==============================================================================
# Section 1: Argument Parsing & Configuration
# ==============================================================================
def parse_args():
    parser = argparse.ArgumentParser(description="ViuAI Sarus-500M DPO Training")
    parser.add_argument("--version", type=str, default="v2", help="DPO version tag (e.g. v2)")
    parser.add_argument("--sft_version", type=str, default="v1", help="Base SFT version tag")
    parser.add_argument("--epochs", type=int, default=2, help="Number of DPO epochs (typically 1-3)")
    parser.add_argument("--batch_size", type=int, default=4, help="Per-device micro batch size")
    parser.add_argument("--grad_accum", type=int, default=4, help="Gradient accumulation steps")
    parser.add_argument("--learning_rate", type=float, default=8.0e-6, help="Peak learning rate for calibrated DPO")
    parser.add_argument("--min_lr", type=float, default=1.0e-6, help="Minimum learning rate")
    parser.add_argument("--beta", type=float, default=0.1, help="DPO temperature beta (0.05 - 0.2)")
    parser.add_argument("--sft_weight", type=float, default=0.15, help="Calibrated auxiliary SFT loss (eliminates attractor loops)")
    parser.add_argument("--max_seq_len", type=int, default=1024, help="Maximum sequence length")
    parser.add_argument("--eval_interval", type=int, default=50, help="Steps between validation evals")
    parser.add_argument("--push_to_hf", action="store_true", help="Auto push checkpoint to HF Hub")
    parser.add_argument("--hf_token", type=str, default="", help="Hugging Face write token")
    return parser.parse_args()

args = parse_args()
# SECURITY: HF token from --hf_token or env only. Never hardcode (it leaks via Hub on every push).
HF_TOKEN = args.hf_token or os.environ.get("HF_TOKEN")
if HF_TOKEN:
    os.environ["HF_TOKEN"] = HF_TOKEN
else:
    print("⚠️ No HF_TOKEN found (checked --hf_token and env). Private Hub downloads/uploads will fail; export HF_TOKEN first.")

IS_DDP = "RANK" in os.environ and "WORLD_SIZE" in os.environ
if IS_DDP:
    LOCAL_RANK = int(os.environ["LOCAL_RANK"])
    WORLD_SIZE = int(os.environ["WORLD_SIZE"])
    RANK = int(os.environ["RANK"])
    torch.cuda.set_device(LOCAL_RANK)
    dist.init_process_group("nccl")
    DEVICE = torch.device(f"cuda:{LOCAL_RANK}")
    IS_MAIN = (RANK == 0)
else:
    DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    LOCAL_RANK = 0
    WORLD_SIZE = 1
    RANK = 0
    IS_MAIN = True

MODEL_REPO = "ViuAI/ViuAI-500M"
DATA_REPO = "ViuAI/viuai-500m-sft-tokenized"
CKPT_DIR = os.path.join("/workspace" if os.path.exists("/workspace") else ".", "dpo_checkpoints", f"dpo_{args.version}")
if IS_MAIN:
    os.makedirs(CKPT_DIR, exist_ok=True)
    print("=" * 85)
    print(f"🎯 ViuAI Sarus-500M — Direct Preference Optimization (DPO {args.version.upper()})")
    print(f"   Base Policy: SFT {args.sft_version.upper()} | Beta: {args.beta} | LR: {args.learning_rate} | GPUs: {WORLD_SIZE}")
    print("=" * 85)

# ==============================================================================
# Section 2: Imports & Tokenizer
# ==============================================================================
code_dir = os.path.abspath("code")
if code_dir not in sys.path:
    sys.path.insert(0, code_dir)

from config import ViuAIConfig
from model import ViuAI
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download, HfApi

tok_path = "tokenizer/tokenizer.json"
if not os.path.exists(tok_path) and IS_MAIN:
    dl = hf_hub_download(repo_id=MODEL_REPO, filename="tokenizer/tokenizer.json", token=HF_TOKEN)
    os.makedirs(os.path.dirname(tok_path), exist_ok=True)
    shutil.copy(dl, tok_path)

if IS_DDP:
    dist.barrier()

tokenizer = Tokenizer.from_file(tok_path)
_tok_vocab = tokenizer.get_vocab()
EOT_ID = _tok_vocab.get("<|endofturn|>", 64002)
# EOT-as-pad (0 is a real vocab token and must not be used as pad).
PAD_ID = EOT_ID

# ==============================================================================
# Section 3: DPO Dataset & Collate
# ==============================================================================
class DPODataset(Dataset):
    def __init__(self, pairs: List[Dict[str, str]], tokenizer: Tokenizer, max_len: int = 1024):
        self.pairs = pairs
        self.tokenizer = tokenizer
        self.max_len = max_len

    def __len__(self):
        return len(self.pairs)

    def _encode_sequence(self, prompt: str, response: str) -> Tuple[List[int], List[int]]:
        prompt_str = f"<|user|>\n{prompt}<|endofturn|>\n<|assistant|>\n"
        prompt_ids = self.tokenizer.encode(prompt_str).ids

        resp_str = f"{response}<|endofturn|>\n"
        resp_ids = self.tokenizer.encode(resp_str).ids

        input_ids = (prompt_ids + resp_ids)[:self.max_len]
        mask = ([0] * len(prompt_ids) + [1] * len(resp_ids))[:self.max_len]
        return input_ids, mask

    def __getitem__(self, idx):
        item = self.pairs[idx]
        prompt = item["prompt"]
        chosen_resp = item["chosen"]
        rejected_resp = item["rejected"]

        chosen_ids, chosen_mask = self._encode_sequence(prompt, chosen_resp)
        rejected_ids, rejected_mask = self._encode_sequence(prompt, rejected_resp)

        return {
            "chosen_input_ids": chosen_ids,
            "chosen_mask": chosen_mask,
            "rejected_input_ids": rejected_ids,
            "rejected_mask": rejected_mask,
            "prompt": prompt,
            "chosen": chosen_resp,
            "rejected": rejected_resp
        }

def dpo_collate_fn(batch):
    def pad_tensors(sequences, pad_val=None):
        if pad_val is None:
            pad_val = PAD_ID
        max_len = max(len(seq) for seq in sequences)
        padded = torch.full((len(sequences), max_len), pad_val, dtype=torch.long)
        for i, seq in enumerate(sequences):
            padded[i, :len(seq)] = torch.tensor(seq, dtype=torch.long)
        return padded

    chosen_ids = pad_tensors([b["chosen_input_ids"] for b in batch], PAD_ID)
    chosen_masks = pad_tensors([b["chosen_mask"] for b in batch], 0)
    rejected_ids = pad_tensors([b["rejected_input_ids"] for b in batch], PAD_ID)
    rejected_masks = pad_tensors([b["rejected_mask"] for b in batch], 0)

    return {
        "chosen_ids": chosen_ids,
        "chosen_mask": chosen_masks,
        "rejected_ids": rejected_ids,
        "rejected_mask": rejected_masks,
        "raw": batch
    }

def load_dpo_data():
    local_train = os.path.join("data", f"dpo_{args.version}", "dpo_train.json")
    local_val = os.path.join("data", f"dpo_{args.version}", "dpo_val.json")

    if not os.path.exists(local_train) and IS_MAIN:
        print("  ⬇️ Fetching DPO pairs from Hugging Face Dataset Hub...")
        dl_t = hf_hub_download(repo_id=DATA_REPO, filename=f"dpo_{args.version}/dpo_train.json", repo_type="dataset", token=HF_TOKEN)
        dl_v = hf_hub_download(repo_id=DATA_REPO, filename=f"dpo_{args.version}/dpo_val.json", repo_type="dataset", token=HF_TOKEN)
        os.makedirs(os.path.dirname(local_train), exist_ok=True)
        shutil.copy(dl_t, local_train)
        shutil.copy(dl_v, local_val)

    if IS_DDP:
        dist.barrier()

    with open(local_train, "r", encoding="utf-8") as f:
        train_pairs = json.load(f)
    with open(local_val, "r", encoding="utf-8") as f:
        val_pairs = json.load(f)

    if IS_MAIN:
        print(f"  ✅ DPO Dataset Loaded: {len(train_pairs):,} Train pairs | {len(val_pairs):,} Val pairs.")
    return train_pairs, val_pairs

# ==============================================================================
# Section 4: Log Probability Computation & DPO Loss
# ==============================================================================
def get_batch_logps(logits: torch.Tensor, labels: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    shift_logits = logits[:, :-1, :].contiguous()
    shift_labels = labels[:, 1:].contiguous()
    shift_mask = mask[:, 1:].contiguous().float()

    log_probs = F.log_softmax(shift_logits, dim=-1)
    per_token_logps = torch.gather(log_probs, 2, shift_labels.unsqueeze(2)).squeeze(2)

    return (per_token_logps * shift_mask).sum(dim=-1)

def compute_dpo_loss(
    policy_chosen_logps: torch.Tensor,
    policy_rejected_logps: torch.Tensor,
    ref_chosen_logps: torch.Tensor,
    ref_rejected_logps: torch.Tensor,
    beta: float
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    pi_logratios = policy_chosen_logps - policy_rejected_logps
    ref_logratios = ref_chosen_logps - ref_rejected_logps

    logits = beta * (pi_logratios - ref_logratios)
    losses = -F.logsigmoid(logits)

    chosen_rewards = beta * (policy_chosen_logps - ref_chosen_logps).detach()
    rejected_rewards = beta * (policy_rejected_logps - ref_rejected_logps).detach()
    accuracy = (chosen_rewards > rejected_rewards).float().mean()
    margin = (chosen_rewards - rejected_rewards).mean()

    return losses.mean(), accuracy, margin, chosen_rewards.mean()

def compute_sft_loss(logits: torch.Tensor, labels: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    shift_logits = logits[:, :-1, :].contiguous()
    shift_labels = labels[:, 1:].contiguous()
    shift_mask = mask[:, 1:].contiguous().float()

    vocab_size = shift_logits.size(-1)
    ce_loss = F.cross_entropy(
        shift_logits.view(-1, vocab_size),
        shift_labels.view(-1),
        reduction='none'
    ).view(shift_labels.size())

    active_tokens = shift_mask.sum().clamp(min=1.0)
    return (ce_loss * shift_mask).sum() / active_tokens

# ==============================================================================
# Section 5: Model Initialization (Policy + Frozen Reference)
# ==============================================================================
def load_models(device, dtype):
    config = ViuAIConfig(vocab_size=64003, context_length=2048)

    sft_ckpt_paths = [
        f"/workspace/sft_checkpoints/sft_{args.sft_version}/sft_{args.sft_version}_final.pt",
        f"sft_checkpoints/sft_{args.sft_version}/sft_{args.sft_version}_final.pt"
    ]
    ckpt_path = None
    for p in sft_ckpt_paths:
        if os.path.exists(p):
            ckpt_path = p
            break

    if ckpt_path is None:
        if IS_MAIN:
            print(f"  ⬇️ Downloading base SFT checkpoint sft_{args.sft_version}_final.pt from Hugging Face...")
            dl = hf_hub_download(repo_id=MODEL_REPO, filename=f"sft_checkpoints/sft_{args.sft_version}/sft_{args.sft_version}_final.pt", token=HF_TOKEN)
            ckpt_path = sft_ckpt_paths[0]
            os.makedirs(os.path.dirname(ckpt_path), exist_ok=True)
            shutil.copy(dl, ckpt_path)

    if IS_DDP:
        dist.barrier()

    if ckpt_path is None:
        ckpt_path = sft_ckpt_paths[0]

    payload = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    state_dict = payload.get("model", payload.get("model_state_dict", payload))
    cleaned_sd = {k.replace("_orig_mod.", "").replace("module.", ""): v for k, v in state_dict.items()}

    # 1. Policy Model (Trainable)
    policy_model = ViuAI(config)
    policy_model.load_state_dict(cleaned_sd)
    policy_model.to(device=device, dtype=dtype)
    policy_model.train()

    # 2. Reference Model (Frozen)
    ref_model = ViuAI(config)
    ref_model.load_state_dict(cleaned_sd)
    ref_model.to(device=device, dtype=dtype)
    for p in ref_model.parameters():
        p.requires_grad = False
    ref_model.eval()

    if IS_MAIN:
        print(f"  ✅ Policy Model & Frozen Reference Model initialized on {device} ({dtype}).")
    return policy_model, ref_model

# ==============================================================================
# Section 6: Main DPO Training Loop
# ==============================================================================
def main():
    dtype = torch.bfloat16 if (torch.cuda.is_available() and torch.cuda.is_bf16_supported()) else torch.float16 if torch.cuda.is_available() else torch.float32

    policy_model, ref_model = load_models(DEVICE, dtype)
    train_pairs, val_pairs = load_dpo_data()

    train_dataset = DPODataset(train_pairs, tokenizer, max_len=args.max_seq_len)
    val_dataset = DPODataset(val_pairs, tokenizer, max_len=args.max_seq_len)

    train_sampler = DistributedSampler(train_dataset, num_replicas=WORLD_SIZE, rank=RANK, shuffle=True) if IS_DDP else None
    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None), sampler=train_sampler, collate_fn=dpo_collate_fn)
    val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, collate_fn=dpo_collate_fn)

    total_steps = (len(train_loader) // args.grad_accum) * args.epochs
    optimizer = torch.optim.AdamW(policy_model.parameters(), lr=args.learning_rate, weight_decay=0.01, betas=(0.9, 0.95))

    def get_lr(step):
        if step < int(0.05 * total_steps):
            return args.learning_rate * (step + 1) / int(0.05 * total_steps)
        progress = (step - int(0.05 * total_steps)) / max(1, total_steps - int(0.05 * total_steps))
        return args.min_lr + 0.5 * (args.learning_rate - args.min_lr) * (1.0 + math.cos(math.pi * progress))

    if IS_MAIN:
        print("\n" + "=" * 85)
        print(f"🏁 STARTING DPO {args.version.upper()} TRAINING ({args.epochs} Epochs | {total_steps} Steps)")
        print("=" * 85)

    global_step = 0
    best_val_acc = 0.0
    start_time = time.time()

    def evaluate():
        policy_model.eval()
        val_losses, val_accs, val_margins = [], [], []

        with torch.no_grad():
            for batch in val_loader:
                c_ids, c_mask = batch["chosen_ids"].to(DEVICE), batch["chosen_mask"].to(DEVICE)
                r_ids, r_mask = batch["rejected_ids"].to(DEVICE), batch["rejected_mask"].to(DEVICE)

                with torch.autocast(device_type=DEVICE.type, dtype=dtype):
                    p_c_logits = policy_model(c_ids, pad_id=PAD_ID)[0]
                    p_r_logits = policy_model(r_ids, pad_id=PAD_ID)[0]
                    ref_c_logits = ref_model(c_ids, pad_id=PAD_ID)[0]
                    ref_r_logits = ref_model(r_ids, pad_id=PAD_ID)[0]

                    p_c_logps = get_batch_logps(p_c_logits, c_ids, c_mask)
                    p_r_logps = get_batch_logps(p_r_logits, r_ids, r_mask)
                    ref_c_logps = get_batch_logps(ref_c_logits, c_ids, c_mask)
                    ref_r_logps = get_batch_logps(ref_r_logits, r_ids, r_mask)

                    loss, acc, margin, _ = compute_dpo_loss(p_c_logps, p_r_logps, ref_c_logps, ref_r_logps, args.beta)

                val_losses.append(loss.item())
                val_accs.append(acc.item())
                val_margins.append(margin.item())

        policy_model.train()
        mean_loss = sum(val_losses) / max(1, len(val_losses))
        mean_acc = sum(val_accs) / max(1, len(val_accs))
        mean_margin = sum(val_margins) / max(1, len(val_margins))
        return mean_loss, mean_acc, mean_margin

    for epoch in range(1, args.epochs + 1):
        if IS_DDP and train_sampler is not None:
            train_sampler.set_epoch(epoch)

        accum_loss, accum_dpo, accum_sft, accum_acc, accum_margin = 0.0, 0.0, 0.0, 0.0, 0.0
        optimizer.zero_grad(set_to_none=True)

        for step, batch in enumerate(train_loader):
            c_ids, c_mask = batch["chosen_ids"].to(DEVICE), batch["chosen_mask"].to(DEVICE)
            r_ids, r_mask = batch["rejected_ids"].to(DEVICE), batch["rejected_mask"].to(DEVICE)

            with torch.autocast(device_type=DEVICE.type, dtype=dtype):
                p_c_logits = policy_model(c_ids, pad_id=PAD_ID)[0]
                p_r_logits = policy_model(r_ids, pad_id=PAD_ID)[0]

                with torch.no_grad():
                    ref_c_logits = ref_model(c_ids, pad_id=PAD_ID)[0]
                    ref_r_logits = ref_model(r_ids, pad_id=PAD_ID)[0]

                p_c_logps = get_batch_logps(p_c_logits, c_ids, c_mask)
                p_r_logps = get_batch_logps(p_r_logits, r_ids, r_mask)
                ref_c_logps = get_batch_logps(ref_c_logits, c_ids, c_mask)
                ref_r_logps = get_batch_logps(ref_r_logits, r_ids, r_mask)

                dpo_loss, acc, margin, _ = compute_dpo_loss(p_c_logps, p_r_logps, ref_c_logps, ref_r_logps, args.beta)
                sft_loss = compute_sft_loss(p_c_logits, c_ids, c_mask) if args.sft_weight > 0 else torch.tensor(0.0, device=DEVICE)
                total_loss = dpo_loss + args.sft_weight * sft_loss
                scaled_loss = total_loss / args.grad_accum

            scaled_loss.backward()
            accum_loss += total_loss.item() / args.grad_accum
            accum_dpo += dpo_loss.item() / args.grad_accum
            accum_sft += sft_loss.item() / args.grad_accum
            accum_acc += acc.item() / args.grad_accum
            accum_margin += margin.item() / args.grad_accum

            if (step + 1) % args.grad_accum == 0 or (step + 1) == len(train_loader):
                torch.nn.utils.clip_grad_norm_(policy_model.parameters(), 1.0)
                lr = get_lr(global_step)
                for param_group in optimizer.param_groups:
                    param_group['lr'] = lr
                optimizer.step()
                optimizer.zero_grad(set_to_none=True)
                global_step += 1

                if global_step % 10 == 0 and IS_MAIN:
                    print(f"Step {global_step:4d}/{total_steps} | Epoch {epoch} | Loss: {accum_loss:.4f} (DPO: {accum_dpo:.3f}, SFT: {accum_sft:.3f}) | Margin: {accum_margin:+.3f} | Pair Acc: {accum_acc*100:5.1f}% | LR: {lr:.2e}")

                if (global_step % args.eval_interval == 0 or global_step == total_steps) and IS_MAIN:
                    val_loss, val_acc, val_margin = evaluate()
                    print("\n" + "-" * 85)
                    print(f"★ [Eval @ Step {global_step}] Val DPO Loss: {val_loss:.4f} | Reward Margin: {val_margin:+.3f} | Pairwise Accuracy: {val_acc*100:.2f}%")
                    print("-" * 85 + "\n")

                    if val_acc >= best_val_acc:
                        best_val_acc = val_acc
                        save_p = os.path.join(CKPT_DIR, f"dpo_{args.version}_final.pt")
                        state = policy_model.module.state_dict() if hasattr(policy_model, "module") else policy_model.state_dict()
                        torch.save({
                            "model_state_dict": state,
                            "val_accuracy": val_acc,
                            "val_loss": val_loss,
                            "step": global_step
                        }, save_p)
                        print(f"  🏆 New Best Pairwise Accuracy ({val_acc*100:.2f}%)! Saved to: {save_p}")

                accum_loss, accum_dpo, accum_sft, accum_acc, accum_margin = 0.0, 0.0, 0.0, 0.0, 0.0

    if IS_MAIN:
        final_save_p = os.path.join(CKPT_DIR, f"dpo_{args.version}_final.pt")
        if not os.path.exists(final_save_p):
            state = policy_model.module.state_dict() if hasattr(policy_model, "module") else policy_model.state_dict()
            torch.save({
                "model_state_dict": state,
                "final_step": global_step,
                "version": args.version
            }, final_save_p)

        print("\n" + "=" * 85)
        print(f"🎉 DPO {args.version.upper()} ALIGNMENT COMPLETED!")
        print(f"   • Final Model Saved: {final_save_p}")
        print(f"   • Total Time: {round((time.time()-start_time)/60, 2)} minutes")
        print("=" * 85)

        if args.push_to_hf:
            print(f"\n🚀 Uploading Final DPO Model to Hugging Face Hub ({MODEL_REPO})...")
            api = HfApi(token=HF_TOKEN)
            api.upload_file(
                path_or_fileobj=final_save_p,
                path_in_repo=f"dpo_checkpoints/dpo_{args.version}/dpo_{args.version}_final.pt",
                repo_id=MODEL_REPO,
                repo_type="model",
                commit_message=f"Upload final DPO {args.version} aligned model"
            )
            print("✅ Successfully uploaded to Hugging Face!")

    if IS_DDP:
        dist.destroy_process_group()

if __name__ == "__main__":
    main()