poimen-memory/verl-training-harness.py

#!/usr/bin/env python3
"""
M5.5 — verl Training Harness

Trains the memory controller on exported trajectories using verl.
Supports both trajectory-level and turn-level policy gradient.

Usage:
    python verl-training-harness.py \\
        --corpus-path corpus/trajectories.jsonl \\
        --output-dir ./checkpoints \\
        --num-epochs 3 \\
        --batch-size 8

Prerequisites:
    - verl installed: pip install verl
    - transformers, peft, trl installed
    - CUDA/GPU available
    - Corpus exported from M5.3
"""

import argparse
import json
import logging
from pathlib import Path
from typing import List, Dict, Any

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import LoraConfig, get_peft_model
from torch.utils.data import Dataset, DataLoader

logger = logging.getLogger(__name__)


class TrajectoryDataset(Dataset):
    """Loads JSONL trajectory format for training."""
    
    def __init__(self, corpus_path: str, tokenizer=None):
        self.corpus_path = Path(corpus_path)
        self.trajectories = []
        self.tokenizer = tokenizer
        self._load_trajectories()
    
    def _load_trajectories(self):
        """Load trajectories from JSONL."""
        with open(self.corpus_path) as f:
            for line in f:
                traj = json.loads(line)
                self.trajectories.append(traj)
        logger.info(f"Loaded {len(self.trajectories)} trajectories")
    
    def __len__(self):
        return len(self.trajectories)
    
    def __getitem__(self, idx: int) -> Dict[str, Any]:
        """Return a trajectory with tokenized prompts/responses."""
        traj = self.trajectories[idx]
        
        turns = traj.get("turns", [])
        rewards = {
            "r_update": [],
            "r_exit": traj.get("r_exit", 0.0),
            "r_format": traj.get("r_format", 1.0),
            "r_outcome": traj.get("r_outcome"),
        }
        
        # Collect turn rewards
        for turn in turns:
            rewards["r_update"].append(turn.get("r_update", 0))
        
        return {
            "trajectory_id": traj.get("trajectory_id"),
            "turns": turns,
            "rewards": rewards,
            "num_turns": len(turns),
        }


def build_model(model_name: str, lora_rank: int = 32):
    """Build base model + LoRA adapter."""
    
    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    # Load base model
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        device_map="auto",
    )
    
    # Configure LoRA
    lora_config = LoraConfig(
        r=lora_rank,
        lora_alpha=32,
        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM",
    )
    
    # Apply LoRA
    model = get_peft_model(model, lora_config)
    
    return model, tokenizer


class PolicyGradientTrainer:
    """Trains with α-blended trajectory + turn level rewards."""
    
    def __init__(
        self,
        model,
        tokenizer,
        learning_rate: float = 5e-5,
        trajectory_weight: float = 0.9,
        turn_weight: float = 0.1,
    ):
        self.model = model
        self.tokenizer = tokenizer
        self.optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=learning_rate,
        )
        self.trajectory_weight = trajectory_weight
        self.turn_weight = turn_weight
    
    def compute_trajectory_loss(
        self,
        turns: List[Dict],
        rewards: Dict,
    ) -> torch.Tensor:
        """Compute trajectory-level loss (α term)."""
        
        r_exit = torch.tensor(rewards["r_exit"], dtype=torch.float32)
        r_format = torch.tensor(rewards["r_format"], dtype=torch.float32)
        
        # Trajectory reward: weighted combination
        traj_reward = 0.7 * r_exit + 0.3 * r_format
        
        return traj_reward
    
    def compute_turn_loss(
        self,
        turns: List[Dict],
        rewards: Dict,
    ) -> torch.Tensor:
        """Compute per-turn loss (1-α term)."""
        
        r_updates = torch.tensor(
            rewards["r_update"],
            dtype=torch.float32,
        )
        
        # Average turn reward
        turn_loss = -r_updates.mean()  # Negative because we minimize loss
        
        return turn_loss
    
    def train_step(self, batch: Dict[str, Any]) -> float:
        """Single training step on a trajectory."""
        
        turns = batch["turns"]
        rewards = batch["rewards"]
        
        # Compute loss components
        traj_loss = self.compute_trajectory_loss(turns, rewards)
        turn_loss = self.compute_turn_loss(turns, rewards)
        
        # Blend losses
        loss = (
            self.trajectory_weight * traj_loss +
            self.turn_weight * turn_loss
        )
        
        # Backward pass
        self.optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
        self.optimizer.step()
        
        return loss.item()


def train(
    corpus_path: str,
    model_name: str = "Qwen/Qwen2.5-3B-Instruct",
    output_dir: str = "./checkpoints",
    num_epochs: int = 3,
    batch_size: int = 8,
    lora_rank: int = 32,
    learning_rate: float = 5e-5,
):
    """Main training loop."""
    
    # Setup
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    logger.info(f"Building model: {model_name}")
    model, tokenizer = build_model(model_name, lora_rank)
    
    logger.info(f"Loading corpus: {corpus_path}")
    dataset = TrajectoryDataset(corpus_path, tokenizer)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    
    logger.info(f"Initializing trainer with lr={learning_rate}")
    trainer = PolicyGradientTrainer(
        model,
        tokenizer,
        learning_rate=learning_rate,
    )
    
    # Training loop
    total_loss = 0.0
    total_steps = 0
    
    for epoch in range(num_epochs):
        logger.info(f"Epoch {epoch+1}/{num_epochs}")
        epoch_loss = 0.0
        
        for step, batch in enumerate(dataloader):
            loss = trainer.train_step(batch)
            epoch_loss += loss
            total_loss += loss
            total_steps += 1
            
            if step % 10 == 0:
                logger.info(f"  Step {step}: loss={loss:.4f}")
        
        avg_epoch_loss = epoch_loss / len(dataloader)
        logger.info(f"Epoch {epoch+1} avg loss: {avg_epoch_loss:.4f}")
        
        # Save checkpoint
        checkpoint_path = output_path / f"memory-v{epoch+1}"
        checkpoint_path.mkdir(parents=True, exist_ok=True)
        model.save_pretrained(checkpoint_path / "adapter")
        tokenizer.save_pretrained(checkpoint_path / "tokenizer")
        logger.info(f"Saved checkpoint: {checkpoint_path}")
    
    # Final summary
    avg_loss = total_loss / total_steps
    logger.info(f"Training complete!")
    logger.info(f"Total steps: {total_steps}")
    logger.info(f"Average loss: {avg_loss:.4f}")
    logger.info(f"Best checkpoint: {output_path / f'memory-v{num_epochs}'}")
    
    return {
        "final_loss": avg_loss,
        "steps_trained": total_steps,
        "epochs": num_epochs,
        "checkpoint_path": str(output_path / f"memory-v{num_epochs}"),
    }


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train memory controller with verl")
    parser.add_argument("--corpus-path", required=True, help="Path to JSONL corpus")
    parser.add_argument("--output-dir", default="./checkpoints", help="Output directory")
    parser.add_argument("--model", default="Qwen/Qwen2.5-3B-Instruct")
    parser.add_argument("--num-epochs", type=int, default=3)
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--lora-rank", type=int, default=32)
    parser.add_argument("--learning-rate", type=float, default=5e-5)
    
    args = parser.parse_args()
    
    logging.basicConfig(level=logging.INFO)
    
    result = train(
        corpus_path=args.corpus_path,
        model_name=args.model,
        output_dir=args.output_dir,
        num_epochs=args.num_epochs,
        batch_size=args.batch_size,
        lora_rank=args.lora_rank,
        learning_rate=args.learning_rate,
    )
    
    print(json.dumps(result, indent=2))