#!/usr/bin/env python3 """ DPO (Direct Preference Optimization) Training Script for Fusion-LLM Implements DPO alignment training using preference pairs (chosen vs rejected). Based on Rafailov et al. (2023) "Direct Preference Optimization". Usage: python train/dpo_finetune.py \ --model_path output/mini_model \ --data_path data/preference_data.json \ --output_path output/dpo_model \ --epochs 3 \ --lr 1e-6 \ --beta 0.1 Author: Zhu Zizhan Project: Fusion-LLM License: Apache 2.0 """ import json import math import os import sys from dataclasses import dataclass, field from pathlib import Path from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F # Add project root to path PROJECT_ROOT = Path(__file__).parent.parent sys.path.insert(0, str(PROJECT_ROOT)) from models.fusion_model import FusionModel, FusionConfig @dataclass class DPOConfig: """DPO training configuration.""" model_path: str = "output/mini_model" data_path: str = "data/preference_data.json" output_path: str = "output/dpo_model" beta: float = 0.1 # DPO temperature parameter lr: float = 1e-6 # Learning rate epochs: int = 3 batch_size: int = 2 max_seq_len: int = 256 warmup_steps: int = 50 gradient_accumulation: int = 4 save_every: int = 100 # Save checkpoint every N steps config_overrides: str = "" # JSON config overrides class DPOTrainer: """ DPO Trainer for Fusion models. Minimizes the DPO loss: L_DPO = -E[log sigmoid(beta * (log pi(y_w|x)/pi_ref(y_w|x) - log pi(y_l|x)/pi_ref(y_l|x)))] Where: - pi: current policy (model being trained) - pi_ref: reference policy (frozen copy of initial model) - y_w: chosen (preferred) response - y_l: rejected response """ def __init__(self, config: DPOConfig): self.config = config self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.step = 0 def _get_tokenizer(self) -> object: """Get tokenizer with fallback to character-level encoding.""" try: from models.tokenizer import get_tokenizer return get_tokenizer("fusion") except Exception: return None def _tokenize(self, text: str, max_len: int) -> torch.Tensor: """Tokenize text using proper tokenizer, falling back to character-level.""" if self._tokenizer is not None: encoded = self._tokenizer.encode(text, truncation=True, max_length=max_len, padding='max_length') return torch.tensor(encoded, dtype=torch.long) # Fallback: UTF-8 byte-level encoding encoded = list(text.encode('utf-8'))[:max_len] padded = encoded + [0] * (max_len - len(encoded)) return torch.tensor(padded, dtype=torch.long) def _detokenize(self, token_ids: list) -> str: """Convert token IDs back to text.""" if self._tokenizer is not None: return self._tokenizer.decode(token_ids, skip_special_tokens=True) try: return bytes(token_ids).decode('utf-8', errors='replace') except Exception: return "" def load_model(self) -> tuple: """Load model and create reference copy.""" config_path = Path(self.config.model_path) if config_path.joinpath("config.json").exists(): model_config = FusionConfig.from_pretrained(str(config_path)) else: # Use mini config as default model_config = FusionConfig( vocab_size=10000, hidden_size=256, num_hidden_layers=2, num_attention_heads=4, intermediate_size=512, block_size=64, latent_dim=16, max_position_embeddings=256, ) # Apply overrides if self.config.config_overrides: overrides = json.loads(self.config.config_overrides) for k, v in overrides.items(): setattr(model_config, k, v) # Policy model (will be trained) policy = FusionModel(model_config) # Load weights if available weight_path = config_path / "final_model.pth" if weight_path.exists(): state_dict = torch.load(weight_path, map_location="cpu", weights_only=True) policy.load_state_dict(state_dict, strict=False) print(f"Loaded weights from {weight_path}") policy = policy.to(self.device) # Reference model (frozen copy) ref_policy = FusionModel(model_config) ref_policy.load_state_dict(policy.state_dict()) ref_policy = ref_policy.to(self.device) ref_policy.eval() for p in ref_policy.parameters(): p.requires_grad = False # Initialize tokenizer self._tokenizer = self._get_tokenizer() if self._tokenizer is not None: print(f"Using tokenizer: vocab_size={self._tokenizer.vocab_size}") else: print("Warning: No tokenizer available, using UTF-8 byte-level fallback") return policy, ref_policy def load_data(self) -> list: """Load preference data.""" data_path = Path(self.config.data_path) if not data_path.exists(): print(f"No preference data at {data_path}, generating synthetic data...") return self._generate_synthetic_data() with open(data_path, 'r', encoding='utf-8') as f: data = json.load(f) print(f"Loaded {len(data)} preference pairs from {data_path}") return data def _generate_synthetic_data(self) -> list: """Generate synthetic preference data for testing.""" pairs = [] templates = [ {"prompt": "What is machine learning", "chosen": "Machine learning is a subset of AI that enables systems to learn from data.", "rejected": "ML is just statistics."}, {"prompt": "Explain neural networks", "chosen": "Neural networks are computing systems inspired by biological neural networks, consisting of interconnected nodes that process information.", "rejected": "They are like brains."}, {"prompt": "What is deep learning", "chosen": "Deep learning uses multi-layered neural networks to automatically learn hierarchical representations from data.", "rejected": "It is deep ML."}, {"prompt": "What is Python", "chosen": "Python is a high-level, interpreted programming language known for its readability and extensive ecosystem of libraries.", "rejected": "A snake."}, {"prompt": "Explain gradient descent", "chosen": "Gradient descent is an optimization algorithm that iteratively moves parameters in the direction of steepest decrease of the loss function.", "rejected": "Going downhill."}, ] for t in templates: pairs.append({ "prompt": t["prompt"], "chosen": t["chosen"], "rejected": t["rejected"], }) # Save for reuse out_path = Path("data/preference_data.json") out_path.parent.mkdir(parents=True, exist_ok=True) with open(out_path, 'w', encoding='utf-8') as f: json.dump(pairs, f, ensure_ascii=False, indent=2) print(f"Generated {len(pairs)} synthetic pairs, saved to {out_path}") return pairs def _get_log_probs(self, model: nn.Module, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: """Compute log probabilities of sequences under model.""" outputs = model(input_ids=input_ids, attention_mask=attention_mask) logits = outputs.logits if hasattr(outputs, 'logits') else outputs['logits'] # Shift for next-token prediction shift_logits = logits[:, :-1, :].contiguous() shift_labels = input_ids[:, 1:].contiguous() # Log softmax log_probs = F.log_softmax(shift_logits, dim=-1) # Gather log probs for actual tokens per_token_log_probs = log_probs.gather(2, shift_labels.unsqueeze(2)).squeeze(2) # Mask padding mask = (shift_labels != 0).float() return (per_token_log_probs * mask).sum(dim=1) def dpo_loss( self, policy_chosen_logps: torch.Tensor, policy_rejected_logps: torch.Tensor, ref_chosen_logps: torch.Tensor, ref_rejected_logps: torch.Tensor, ) -> torch.Tensor: """Compute DPO loss.""" chosen_rewards = self.config.beta * (policy_chosen_logps - ref_chosen_logps) rejected_rewards = self.config.beta * (policy_rejected_logps - ref_rejected_logps) loss = -F.logsigmoid(chosen_rewards - rejected_rewards).mean() return loss def train(self): """Run DPO training.""" policy, ref_policy = self.load_model() data = self.load_data() optimizer = torch.optim.AdamW(policy.parameters(), lr=self.config.lr) scheduler = torch.optim.lr_scheduler.LinearLR( optimizer, start_factor=0.1, total_iters=self.config.warmup_steps ) output_dir = Path(self.config.output_path) output_dir.mkdir(parents=True, exist_ok=True) print(f"\n[DPO] Starting training:") print(f" Model: {self.config.model_path}") print(f" Data: {len(data)} pairs") print(f" Beta: {self.config.beta}") print(f" LR: {self.config.lr}") print(f" Epochs: {self.config.epochs}") print(f" Device: {self.device}") policy.train() global_step = 0 for epoch in range(self.config.epochs): total_loss = 0.0 num_batches = 0 # Shuffle data each epoch indices = torch.randperm(len(data)) for i in range(0, len(data), self.config.batch_size): batch_indices = indices[i:i + self.config.batch_size] chosen_ids = [] rejected_ids = [] for idx in batch_indices: item = data[idx.item()] prompt = item['prompt'] chosen_ids.append(self._tokenize(prompt + " " + item['chosen'], self.config.max_seq_len)) rejected_ids.append(self._tokenize(prompt + " " + item['rejected'], self.config.max_seq_len)) chosen_ids = torch.stack(chosen_ids).to(self.device) rejected_ids = torch.stack(rejected_ids).to(self.device) chosen_mask = (chosen_ids != 0).float() rejected_mask = (rejected_ids != 0).float() # Policy log probs policy_chosen_logps = self._get_log_probs(policy, chosen_ids, chosen_mask) policy_rejected_logps = self._get_log_probs(policy, rejected_ids, rejected_mask) # Reference log probs (no grad) with torch.no_grad(): ref_chosen_logps = self._get_log_probs(ref_policy, chosen_ids, chosen_mask) ref_rejected_logps = self._get_log_probs(ref_policy, rejected_ids, rejected_mask) # DPO loss loss = self.dpo_loss( policy_chosen_logps, policy_rejected_logps, ref_chosen_logps, ref_rejected_logps, ) # Gradient accumulation loss = loss / self.config.gradient_accumulation loss.backward() if (global_step + 1) % self.config.gradient_accumulation == 0: torch.nn.utils.clip_grad_norm_(policy.parameters(), 1.0) optimizer.step() scheduler.step() optimizer.zero_grad() total_loss += loss.item() * self.config.gradient_accumulation num_batches += 1 global_step += 1 # Save checkpoint if self.config.save_every > 0 and global_step % self.config.save_every == 0: ckpt_path = output_dir / f"checkpoint-step-{global_step}.pth" torch.save(policy.state_dict(), ckpt_path) avg_loss = total_loss / max(num_batches, 1) print(f"[DPO] Epoch {epoch+1}/{self.config.epochs} - Loss: {avg_loss:.4f} - Steps: {global_step}") # Save final model final_path = output_dir / "dpo_model.pth" torch.save(policy.state_dict(), final_path) policy.config.save_pretrained(str(output_dir)) print(f"\n[DPO] Training complete! Saved to {output_dir}") return policy if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="DPO Training for Fusion-LLM") parser.add_argument("--model_path", type=str, default="output/mini_model") parser.add_argument("--data_path", type=str, default="data/preference_data.json") parser.add_argument("--output_path", type=str, default="output/dpo_model") parser.add_argument("--beta", type=float, default=0.1) parser.add_argument("--lr", type=float, default=1e-6) parser.add_argument("--epochs", type=int, default=3) parser.add_argument("--batch_size", type=int, default=2) parser.add_argument("--max_seq_len", type=int, default=256) args = parser.parse_args() config = DPOConfig( model_path=args.model_path, data_path=args.data_path, output_path=args.output_path, beta=args.beta, lr=args.lr, epochs=args.epochs, batch_size=args.batch_size, max_seq_len=args.max_seq_len, ) trainer = DPOTrainer(config) trainer.train()