Upload scripts/train_alizee_v2_stage1_sft.py with huggingface_hub

scripts/train_alizee_v2_stage1_sft.py

@@ -0,0 +1,330 @@
+#!/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 1: Reasoning Distillation via SFT for Alizee-Coder-Devstral-2-Small
+
+Training stmasson/alizee-coder-devstral-1-small on nvidia/OpenCodeReasoning (736K samples)
+with 85% reasoning traces + 15% coding capability preservation from bigcode/starcoderdata.
+
+Key features:
+- QLoRA (r=64, alpha=128) for memory-efficient training
+- 32K context window support
+- Gradient checkpointing + Flash Attention 2
+- Automatic data mixing and formatting
+- Trackio monitoring
+
+Based on NVIDIA's research: performance improves linearly 25K->736K samples
+"""
+
+import os
+import random
+import trackio
+from datasets import load_dataset, concatenate_datasets, Dataset
+from peft import LoraConfig, prepare_model_for_kbit_training, get_peft_model
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from trl import SFTTrainer, SFTConfig
+
+# Configuration
+MODEL_NAME = "stmasson/alizee-coder-devstral-1-small"
+OUTPUT_REPO = "stmasson/alizee-coder-devstral-2-small-stage1"
+FINAL_REPO = "stmasson/alizee-coder-devstral-2-small"
+
+# Training hyperparameters (from user spec)
+LEARNING_RATE = 5e-5
+EFFECTIVE_BATCH_SIZE = 256
+PER_DEVICE_BATCH = 1
+GRADIENT_ACCUMULATION = EFFECTIVE_BATCH_SIZE // PER_DEVICE_BATCH
+MAX_SEQ_LENGTH = 32768
+NUM_EPOCHS = 2
+WARMUP_RATIO = 0.05
+
+# Data mixing ratio
+REASONING_RATIO = 0.85
+CODING_RATIO = 0.15
+
+print("=" * 60)
+print("Stage 1: Reasoning Distillation via SFT")
+print("=" * 60)
+print(f"Base model: {MODEL_NAME}")
+print(f"Output: {OUTPUT_REPO}")
+print(f"Data mix: {REASONING_RATIO*100}% reasoning + {CODING_RATIO*100}% coding")
+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 = "right"
+
+# 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
+print("\n🔄 Loading 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 (r=64, alpha=128 as specified)
+print("\n🎯 Configuring LoRA adapters...")
+lora_config = LoraConfig(
+    r=64,
+    lora_alpha=128,
+    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 and prepare datasets
+print("\n📦 Loading datasets...")
+
+# 1. OpenCodeReasoning (reasoning traces)
+print("   Loading nvidia/OpenCodeReasoning split_0...")
+ocr_split0 = load_dataset("nvidia/OpenCodeReasoning", "split_0", split="train")
+print(f"   -> split_0: {len(ocr_split0)} samples")
+
+print("   Loading nvidia/OpenCodeReasoning split_1...")
+ocr_split1 = load_dataset("nvidia/OpenCodeReasoning", "split_1", split="train")
+print(f"   -> split_1: {len(ocr_split1)} samples")
+
+# Combine OpenCodeReasoning splits
+ocr_full = concatenate_datasets([ocr_split0, ocr_split1])
+print(f"   Total OpenCodeReasoning: {len(ocr_full)} samples")
+
+# 2. Coding capability preservation dataset
+print("   Loading bigcode/starcoderdata (python subset)...")
+# Load a subset of starcoderdata for coding preservation
+coding_ds = load_dataset(
+    "bigcode/starcoderdata",
+    data_dir="python",
+    split="train",
+    streaming=True
+)
+
+# Calculate how many coding samples we need (15% of total)
+total_reasoning = len(ocr_full)
+num_coding_samples = int(total_reasoning * CODING_RATIO / REASONING_RATIO)
+print(f"   Need {num_coding_samples} coding samples for 15% mix")
+
+# Take samples from streaming dataset
+print("   Sampling coding data...")
+coding_samples = []
+for i, sample in enumerate(coding_ds):
+    if i >= num_coding_samples:
+        break
+    coding_samples.append(sample)
+    if i % 50000 == 0 and i > 0:
+        print(f"      Collected {i} coding samples...")
+
+coding_ds_final = Dataset.from_list(coding_samples)
+print(f"   Collected {len(coding_ds_final)} coding samples")
+
+# Format functions for different data sources
+def format_reasoning_sample(example):
+    """Format OpenCodeReasoning sample for instruction tuning.
+
+    OpenCodeReasoning has:
+    - input: problem description
+    - output: reasoning trace / expected output explanation
+    - solution: the actual code
+    """
+    # Create a reasoning-enhanced prompt
+    messages = [
+        {
+            "role": "user",
+            "content": f"Solve the following programming problem. Think through it step by step.\n\n{example['input']}"
+        },
+        {
+            "role": "assistant",
+            "content": f"Let me think through this problem step by step.\n\n{example['output']}\n\nHere's my solution:\n\n```python\n{example['solution']}\n```"
+        }
+    ]
+
+    return {"messages": messages, "source": "reasoning"}
+
+def format_coding_sample(example):
+    """Format starcoderdata sample for capability preservation."""
+    # Extract code content
+    content = example.get("content", "")
+
+    # Create a simple code completion task
+    lines = content.split("\n")
+    if len(lines) > 10:
+        # Split into prompt and completion
+        split_point = len(lines) // 3
+        prompt_code = "\n".join(lines[:split_point])
+        completion_code = "\n".join(lines[split_point:])
+
+        messages = [
+            {
+                "role": "user",
+                "content": f"Complete the following Python code:\n\n```python\n{prompt_code}\n```"
+            },
+            {
+                "role": "assistant",
+                "content": f"```python\n{completion_code}\n```"
+            }
+        ]
+    else:
+        # For short snippets, ask to explain and reproduce
+        messages = [
+            {
+                "role": "user",
+                "content": f"Write Python code that implements the following:\n\n{content[:200]}..."
+            },
+            {
+                "role": "assistant",
+                "content": f"```python\n{content}\n```"
+            }
+        ]
+
+    return {"messages": messages, "source": "coding"}
+
+# Format datasets
+print("\n🔄 Formatting datasets...")
+print("   Formatting reasoning samples...")
+reasoning_formatted = ocr_full.map(
+    format_reasoning_sample,
+    remove_columns=ocr_full.column_names,
+    num_proc=8,
+    desc="Formatting reasoning"
+)
+
+print("   Formatting coding samples...")
+coding_formatted = coding_ds_final.map(
+    format_coding_sample,
+    remove_columns=coding_ds_final.column_names,
+    num_proc=4,
+    desc="Formatting coding"
+)
+
+# Combine and shuffle
+print("\n🔀 Combining and shuffling datasets...")
+combined_dataset = concatenate_datasets([reasoning_formatted, coding_formatted])
+combined_dataset = combined_dataset.shuffle(seed=42)
+print(f"   Total training samples: {len(combined_dataset)}")
+
+# Create train/eval split
+print("   Creating train/eval split...")
+split_dataset = combined_dataset.train_test_split(test_size=0.005, seed=42)
+train_dataset = split_dataset["train"]
+eval_dataset = split_dataset["test"]
+print(f"   Train: {len(train_dataset)}, Eval: {len(eval_dataset)}")
+
+# Training configuration
+print("\n⚙️ Configuring training...")
+training_config = SFTConfig(
+    # Output and Hub settings
+    output_dir="alizee-v2-stage1-sft",
+    push_to_hub=True,
+    hub_model_id=OUTPUT_REPO,
+    hub_strategy="every_save",
+    hub_private_repo=False,
+
+    # 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_seq_length=MAX_SEQ_LENGTH,
+
+    # Optimization
+    warmup_ratio=WARMUP_RATIO,
+    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=500,
+    save_total_limit=3,
+    eval_strategy="steps",
+    eval_steps=500,
+
+    # Monitoring
+    report_to="trackio",
+    project="alizee-coder-v2",
+    run_name="stage1-reasoning-sft",
+
+    # Other settings
+    max_grad_norm=1.0,
+    dataloader_num_workers=4,
+    remove_unused_columns=True,
+    packing=False,  # Disable packing for long sequences
+)
+
+# Initialize trainer
+print("\n🎯 Initializing SFT Trainer...")
+trainer = SFTTrainer(
+    model=model,
+    tokenizer=tokenizer,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    args=training_config,
+)
+
+# Calculate and display training info
+total_steps = (len(train_dataset) // EFFECTIVE_BATCH_SIZE) * NUM_EPOCHS
+print(f"\n📊 Training Configuration Summary:")
+print(f"   Total samples: {len(train_dataset)}")
+print(f"   Effective batch size: {EFFECTIVE_BATCH_SIZE}")
+print(f"   Steps per epoch: {len(train_dataset) // EFFECTIVE_BATCH_SIZE}")
+print(f"   Total steps: {total_steps}")
+print(f"   Epochs: {NUM_EPOCHS}")
+
+# Start training
+print("\n🚀 Starting Stage 1 Reasoning SFT Training...")
+print("   This will take 16-24+ hours on A100-80GB")
+print("   Monitor at: https://huggingface.co/spaces/stmasson/trackio")
+print("=" * 60)
+
+trainer.train()
+
+# Save final model
+print("\n💾 Pushing final model to Hub...")
+trainer.push_to_hub()
+
+# Finish tracking
+trackio.finish()
+
+print("\n" + "=" * 60)
+print("✅ Stage 1 Complete!")
+print(f"   Model saved to: https://huggingface.co/{OUTPUT_REPO}")
+print("   Ready for Stage 2 (optional DPO refresh)")
+print("=" * 60)