train_dpo.py

#!/usr/bin/env python3
"""
Genesis-2.0 — Phase 1: DPO Preference Alignment

Run on RunPod RTX PRO 6000 Blackwell (96 GB).

Steps:
1. Load Qwen3.6-35B-A3B with QLoRA (NF4)
2. Merge existing Genesis-1.0 SFT adapter
3. Attach new all-linear LoRA (Config C: attn + shared + gate, r=32)
4. Load DPO preference pairs
5. Run DPO training
6. Save + upload adapter to HuggingFace Hub

Usage:
    python3 train_dpo.py
"""

import os
import json
import sys
import torch
from typing import Optional

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

# ============================================================
# CONFIG
# ============================================================

class Config:
    # Paths
    base_model = "Qwen/Qwen3.6-35B-A3B"
    sft_adapter = "jacobeen06/Genesis-1.0-SFT-adapter"  # Genesis-1.0 SFT adapter
    dpo_data = "/workspace/dpo_pairs_all.jsonl"  # Will be uploaded
    output_dir = "/workspace/genesis2-dpo"
    hf_repo = "jacobeen06/Genesis-2.0-DPO-adapter"

    # QLoRA
    load_in_4bit = True
    bnb_4bit_quant_type = "nf4"
    bnb_4bit_compute_dtype = torch.bfloat16

    # New all-linear LoRA (Config C)
    lora_r = 32
    lora_alpha = 64
    lora_dropout = 0.0
    # Attention modules
    lora_target_modules = [
        "q_proj", "k_proj", "v_proj", "o_proj",
    ]
    use_rslora = True

    # DPO hyperparameters
    beta = 0.1  # DPO temperature
    learning_rate = 5e-6
    lr_scheduler_type = "cosine"
    warmup_ratio = 0.05
    per_device_train_batch_size = 1
    gradient_accumulation_steps = 4
    num_train_epochs = 1
    max_length = 3072
    max_prompt_length = 2304
    logging_steps = 10
    save_steps = 200
    eval_steps = 200
    save_total_limit = 2
    remove_unused_columns = False

    # DeepSpeed / distributed
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    world_size = int(os.environ.get("WORLD_SIZE", 1))
    deepspeed_config = None  # Use if multi-GPU: "ds_config.json"


# ============================================================
# MODEL LOADING
# ============================================================

def load_model_and_tokenizer(config: Config):
    """Load base model with QLoRA and merge the existing SFT adapter."""
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
    from peft import PeftModel, LoraConfig, get_peft_model

    print(f"[{config.local_rank}] Loading base model: {config.base_model}")

    # Quantization config
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=config.load_in_4bit,
        bnb_4bit_quant_type=config.bnb_4bit_quant_type,
        bnb_4bit_compute_dtype=config.bnb_4bit_compute_dtype,
    )

    model = AutoModelForCausalLM.from_pretrained(
        config.base_model,
        quantization_config=bnb_config,
        device_map="auto" if config.world_size == 1 else {"": config.local_rank},
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
    )

    tokenizer = AutoTokenizer.from_pretrained(
        config.base_model,
        trust_remote_code=True,
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.pad_token_id = tokenizer.eos_token_id

    # ---- Merge Genesis-1.0 SFT adapter ----
    print(f"[{config.local_rank}] Loading Genesis-1.0 SFT adapter...")
    model = PeftModel.from_pretrained(model, config.sft_adapter)
    print(f"[{config.local_rank}] Merging SFT adapter into base...")
    model = model.merge_and_unload()
    print(f"[{config.local_rank}] SFT adapter merged. Model type: {type(model).__name__}")

    # ---- Attach new all-linear LoRA ----
    print(f"[{config.local_rank}] Attaching new all-linear LoRA (r={config.lora_r})...")
    lora_config = LoraConfig(
        r=config.lora_r,
        lora_alpha=config.lora_alpha,
        target_modules=config.lora_target_modules,
        lora_dropout=config.lora_dropout,
        bias="none",
        task_type="CAUSAL_LM",
        use_rslora=config.use_rslora,
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    return model, tokenizer


# ============================================================
# DATA LOADING
# ============================================================

def load_dpo_data(config: Config):
    """Load DPO pairs from JSONL."""
    from datasets import Dataset

    if not os.path.exists(config.dpo_data):
        print(f"[{config.local_rank}] ERROR: {config.dpo_data} not found!")
        print("Upload dpo_pairs_all.jsonl to /workspace/ on the pod first.")
        sys.exit(1)

    pairs = []
    with open(config.dpo_data) as f:
        for line in f:
            line = line.strip()
            if line:
                pairs.append(json.loads(line))

    print(f"[{config.local_rank}] Loaded {len(pairs)} DPO pairs from {config.dpo_data}")

    # Convert to dataset
    dataset = Dataset.from_list(pairs)
    return dataset


# ============================================================
# TRAINING
# ============================================================

def train_dpo(config: Config):
    from transformers import TrainingArguments
    from trl import DPOTrainer, DPOConfig

    model, tokenizer = load_model_and_tokenizer(config)
    dataset = load_dpo_data(config)

    # Split train/eval
    split_dataset = dataset.train_test_split(test_size=0.05, seed=42)
    train_dataset = split_dataset["train"]
    eval_dataset = split_dataset["test"]

    print(f"[{config.local_rank}] Train: {len(train_dataset)}, Eval: {len(eval_dataset)}")

    # Training args (DPOConfig extends TrainingArguments with DPO-specific params)
    training_args = DPOConfig(
        output_dir=config.output_dir,
        per_device_train_batch_size=config.per_device_train_batch_size,
        gradient_accumulation_steps=config.gradient_accumulation_steps,
        learning_rate=config.learning_rate,
        lr_scheduler_type=config.lr_scheduler_type,
        warmup_ratio=config.warmup_ratio,
        num_train_epochs=config.num_train_epochs,
        logging_steps=config.logging_steps,
        save_steps=config.save_steps,
        eval_steps=config.eval_steps,
        eval_strategy="no",  # Skip eval to avoid MoE aux loss OOM
        save_total_limit=config.save_total_limit,
        remove_unused_columns=config.remove_unused_columns,
        bf16=True,
        tf32=True,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        disable_dropout=True,
        router_aux_loss_coef=0.0,
        logging_dir=os.path.join(config.output_dir, "logs"),
        report_to="wandb" if os.environ.get("WANDB_API_KEY") else "none",
        run_name="genesis2-dpo",
        ddp_find_unused_parameters=False if config.world_size > 1 else None,
        dataloader_num_workers=2,
        beta=config.beta,
        max_length=config.max_length,
    )

    # DPO Trainer
    trainer = DPOTrainer(
        model=model,
        ref_model=None,  # Will use a frozen copy internally
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        processing_class=tokenizer,
    )

    # Train
    print(f"[{config.local_rank}] Starting DPO training...")
    trainer.train()

    # Save
    print(f"[{config.local_rank}] Saving model to {config.output_dir}")
    trainer.save_model(config.output_dir)
    tokenizer.save_pretrained(config.output_dir)

    # Upload to HF Hub
    try:
        from huggingface_hub import HfApi
        api = HfApi()
        api.create_repo(config.hf_repo, exist_ok=True)
        api.upload_folder(
            folder_path=config.output_dir,
            repo_id=config.hf_repo,
            commit_message="Genesis-2.0 DPO adapter",
        )
        print(f"[{config.local_rank}] Uploaded to {config.hf_repo}")
    except Exception as e:
        print(f"[{config.local_rank}] Upload failed (non-fatal): {e}")

    print(f"[{config.local_rank}] DONE! Adapter saved to {config.output_dir}")
    return trainer


# ============================================================
# MAIN
# ============================================================

if __name__ == "__main__":
    config = Config()

    print("=" * 60)
    print("Genesis-2.0 — Phase 1: DPO Preference Alignment")
    print("=" * 60)
    print(f"Base model: {config.base_model}")
    print(f"LoRA: r={config.lora_r}, targets={config.lora_target_modules}")
    print(f"DPO data: {config.dpo_data}")
    print(f"Output: {config.output_dir}")
    print()

    train_dpo(config)