# ============================================================================== # šŸš€ 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()