DPO training script

train_dpo.py

@@ -0,0 +1,264 @@
+#!/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
+
+# ============================================================
+# 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",
+        # Shared expert FFN
+        "gate_proj", "up_proj", "down_proj",
+        # Gate/router
+        "gate",
+    ]
+    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 = 4096
+    max_prompt_length = 3072
+    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
+
+    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
+    training_args = TrainingArguments(
+        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="steps" if len(eval_dataset) > 0 else "no",
+        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},
+        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,
+    )
+
+    # 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,
+        tokenizer=tokenizer,
+        beta=config.beta,
+        max_length=config.max_length,
+        max_prompt_length=config.max_prompt_length,
+    )
+
+    # 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)