#!/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)