Upload scripts/train_alizee_v2_stage2_dpo.py with huggingface_hub

scripts/train_alizee_v2_stage2_dpo.py

@@ -0,0 +1,223 @@
+#!/usr/bin/env python3
+# /// script
+# dependencies = [
+#     "trl>=0.17.0",
+#     "peft>=0.14.0",
+#     "transformers>=4.48.0",
+#     "accelerate>=0.35.0",
+#     "bitsandbytes>=0.45.0",
+#     "trackio",
+#     "datasets>=3.0.0",
+#     "flash-attn>=2.5.0",
+# ]
+# ///
+
+"""
+Stage 2: Light DPO Refresh for Alizee-Coder-Devstral-2-Small
+
+Conservative DPO (beta=0.1, lr=5e-6) using CodeUltraFeedback to restore alignment
+after reasoning SFT. This stage is OPTIONAL - run only if evaluation shows
+alignment degradation.
+
+Key settings (from user spec):
+- beta=0.1 (conservative KL penalty)
+- learning_rate=5e-6 (very low to preserve Stage 1 gains)
+"""
+
+import os
+import trackio
+from datasets import load_dataset
+from peft import LoraConfig, prepare_model_for_kbit_training, get_peft_model, PeftModel
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from trl import DPOTrainer, DPOConfig
+
+# Configuration
+MODEL_NAME = "stmasson/alizee-coder-devstral-2-small-stage1"  # Output from Stage 1
+OUTPUT_REPO = "stmasson/alizee-coder-devstral-2-small-stage2"
+
+# DPO hyperparameters (conservative as specified)
+BETA = 0.1  # KL penalty - higher = stay closer to reference
+LEARNING_RATE = 5e-6  # Very low LR for alignment refresh
+EFFECTIVE_BATCH_SIZE = 64
+PER_DEVICE_BATCH = 1
+GRADIENT_ACCUMULATION = EFFECTIVE_BATCH_SIZE // PER_DEVICE_BATCH
+MAX_SEQ_LENGTH = 8192  # Shorter context for DPO
+NUM_EPOCHS = 1
+
+print("=" * 60)
+print("Stage 2: Light DPO Refresh (Optional)")
+print("=" * 60)
+print(f"Base model: {MODEL_NAME}")
+print(f"Output: {OUTPUT_REPO}")
+print(f"Beta (KL penalty): {BETA}")
+print(f"Learning rate: {LEARNING_RATE}")
+print("=" * 60)
+
+# Load tokenizer
+print("\n📝 Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = "left"  # DPO prefers left padding
+
+# QLoRA quantization config
+print("\n⚙️ Configuring 4-bit quantization...")
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype="bfloat16",
+    bnb_4bit_use_double_quant=True,
+)
+
+# Load model (which already has merged LoRA from Stage 1)
+print("\n🔄 Loading Stage 1 model with QLoRA...")
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    quantization_config=bnb_config,
+    device_map="auto",
+    trust_remote_code=True,
+    attn_implementation="flash_attention_2",
+    torch_dtype="auto",
+)
+model = prepare_model_for_kbit_training(model)
+
+# LoRA configuration (smaller for DPO - just alignment refresh)
+print("\n🎯 Configuring LoRA adapters for DPO...")
+lora_config = LoraConfig(
+    r=32,  # Smaller rank for DPO refresh
+    lora_alpha=64,
+    lora_dropout=0.05,
+    bias="none",
+    task_type="CAUSAL_LM",
+    target_modules=[
+        "q_proj", "k_proj", "v_proj", "o_proj",
+        "gate_proj", "up_proj", "down_proj"
+    ],
+)
+
+model = get_peft_model(model, lora_config)
+model.print_trainable_parameters()
+
+# Load CodeUltraFeedback dataset
+print("\n📦 Loading CodeUltraFeedback dataset...")
+dataset = load_dataset("RLHFlow/CodeUltraFeedback-standard", split="train")
+print(f"   Loaded {len(dataset)} preference pairs")
+
+def format_for_dpo(example):
+    """Format CodeUltraFeedback for DPO training.
+
+    CodeUltraFeedback-standard has:
+    - prompt: the coding instruction
+    - chosen: the better response
+    - rejected: the worse response
+    """
+    return {
+        "prompt": example["prompt"],
+        "chosen": example["chosen"],
+        "rejected": example["rejected"],
+    }
+
+# Format dataset
+print("\n🔄 Formatting dataset for DPO...")
+formatted_dataset = dataset.map(
+    format_for_dpo,
+    remove_columns=[col for col in dataset.column_names if col not in ["prompt", "chosen", "rejected"]],
+    num_proc=4,
+)
+
+# Create train/eval split
+print("   Creating train/eval split...")
+split_dataset = formatted_dataset.train_test_split(test_size=0.05, seed=42)
+train_dataset = split_dataset["train"]
+eval_dataset = split_dataset["test"]
+print(f"   Train: {len(train_dataset)}, Eval: {len(eval_dataset)}")
+
+# DPO Training configuration
+print("\n⚙️ Configuring DPO training...")
+training_config = DPOConfig(
+    # Output and Hub settings
+    output_dir="alizee-v2-stage2-dpo",
+    push_to_hub=True,
+    hub_model_id=OUTPUT_REPO,
+    hub_strategy="every_save",
+    hub_private_repo=False,
+
+    # DPO-specific
+    beta=BETA,
+
+    # Training parameters
+    num_train_epochs=NUM_EPOCHS,
+    per_device_train_batch_size=PER_DEVICE_BATCH,
+    per_device_eval_batch_size=PER_DEVICE_BATCH,
+    gradient_accumulation_steps=GRADIENT_ACCUMULATION,
+    learning_rate=LEARNING_RATE,
+    max_length=MAX_SEQ_LENGTH,
+    max_prompt_length=MAX_SEQ_LENGTH // 2,
+
+    # Optimization
+    warmup_ratio=0.1,
+    lr_scheduler_type="cosine",
+    optim="adamw_8bit",
+    bf16=True,
+    gradient_checkpointing=True,
+    gradient_checkpointing_kwargs={"use_reentrant": False},
+
+    # Logging and checkpointing
+    logging_steps=10,
+    save_strategy="steps",
+    save_steps=200,
+    save_total_limit=2,
+    eval_strategy="steps",
+    eval_steps=200,
+
+    # Monitoring
+    report_to="trackio",
+    project="alizee-coder-v2",
+    run_name="stage2-dpo-refresh",
+
+    # Other settings
+    max_grad_norm=1.0,
+    remove_unused_columns=False,
+)
+
+# Initialize trainer
+print("\n🎯 Initializing DPO Trainer...")
+trainer = DPOTrainer(
+    model=model,
+    ref_model=None,  # Use implicit reference (copy of model)
+    tokenizer=tokenizer,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    args=training_config,
+    peft_config=lora_config,
+)
+
+# Calculate and display training info
+total_steps = (len(train_dataset) // EFFECTIVE_BATCH_SIZE) * NUM_EPOCHS
+print(f"\n📊 DPO Training Configuration Summary:")
+print(f"   Total preference pairs: {len(train_dataset)}")
+print(f"   Effective batch size: {EFFECTIVE_BATCH_SIZE}")
+print(f"   Total steps: {total_steps}")
+print(f"   Beta (KL penalty): {BETA}")
+print(f"   Learning rate: {LEARNING_RATE}")
+
+# Start training
+print("\n🚀 Starting Stage 2 DPO Refresh...")
+print("   This should take 2-4 hours on A100-80GB")
+print("   Monitor at: https://huggingface.co/spaces/stmasson/trackio")
+print("=" * 60)
+
+trainer.train()
+
+# Save final model
+print("\n💾 Pushing Stage 2 model to Hub...")
+trainer.push_to_hub()
+
+# Finish tracking
+trackio.finish()
+
+print("\n" + "=" * 60)
+print("✅ Stage 2 Complete!")
+print(f"   Model saved to: https://huggingface.co/{OUTPUT_REPO}")
+print("   Ready for Stage 3 (adapter merging)")
+print("=" * 60)