feat: upgrade finetune.py — LoRA Without Regret + best datasets (tulu-3-sft, OpenThoughts-114k)

ai-ml/hf-finetuning/finetune.py

@@ -1,101 +1,146 @@
 # =============================================================================
-# HuggingFace Fine-Tuning Script — Secure Production Training
+# HuggingFace Fine-Tuning Script — vNext Production Training
 # =============================================================================
-# Uses: TRL SFTTrainer + PEFT LoRA + Trackio monitoring
+# Based on: "LoRA Without Regret" (Schulman et al., 2025)
+#   - LoRA matches full fine-tuning with correct configuration
+#   - Key: all-linear targets + r=256 + LR 2e-4 + batch < 32
+#
+# Datasets (ranked by quality):
+#   PRIMARY:  allenai/tulu-3-sft-mixture (940K examples, 19 sources)
+#   REASONING: open-thoughts/OpenThoughts-114k (CoT traces)
+#   FALLBACK: HuggingFaceH4/ultrachat_200k (200K multi-turn chat)
 # =============================================================================
 
 import os
 import torch
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import Optional, List
 
-from datasets import load_dataset
+from datasets import load_dataset, concatenate_datasets
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
     BitsAndBytesConfig,
-    TrainingArguments,
 )
-from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+from peft import LoraConfig
 from trl import SFTTrainer, SFTConfig
 import trackio
 
 
+# ---------- Dataset Registry ----------
+DATASET_REGISTRY = {
+    "tulu-3-sft": {
+        "name": "allenai/tulu-3-sft-mixture",
+        "split": "train",
+        "format": "messages",  # Already conversational
+        "size": "~940K",
+        "quality": "BEST — 19 curated sources (math, code, IF, safety, science)",
+    },
+    "openthoughts-114k": {
+        "name": "open-thoughts/OpenThoughts-114k",
+        "split": "train",
+        "format": "conversations",  # Needs conversion
+        "size": "~114K",
+        "quality": "EXCELLENT — reasoning CoT traces",
+    },
+    "ultrachat-200k": {
+        "name": "HuggingFaceH4/ultrachat_200k",
+        "split": "train_sft",
+        "format": "messages",
+        "size": "~200K",
+        "quality": "GOOD — multi-turn chat (baseline fallback)",
+    },
+}
+
+
+def convert_openthoughts_to_messages(example):
+    """Convert OpenThoughts conversations format to standard messages."""
+    messages = []
+    if example.get("system"):
+        messages.append({"role": "system", "content": example["system"]})
+    for turn in example["conversations"]:
+        role = "user" if turn["from"] == "user" else "assistant"
+        messages.append({"role": role, "content": turn["value"]})
+    return {"messages": messages}
+
+
+def load_and_prepare_dataset(dataset_key: str, max_samples: Optional[int] = None):
+    """Load and format a dataset from the registry."""
+    info = DATASET_REGISTRY[dataset_key]
+    ds = load_dataset(info["name"], split=info["split"])
+
+    if max_samples:
+        ds = ds.select(range(min(max_samples, len(ds))))
+
+    if dataset_key == "openthoughts-114k":
+        remove_cols = [c for c in ds.column_names if c != "messages"]
+        ds = ds.map(
+            convert_openthoughts_to_messages,
+            remove_columns=remove_cols,
+        )
+
+    return ds
+
+
 @dataclass
 class FinetuneConfig:
-    """Fine-tuning hyperparameters."""
+    """Fine-tuning hyperparameters — vNext (LoRA Without Regret config)."""
     model_name: str = "meta-llama/Llama-3.1-8B-Instruct"
-    dataset_name: str = "HuggingFaceH4/ultrachat_200k"
+    dataset_key: str = "tulu-3-sft"  # Options: tulu-3-sft, openthoughts-114k, ultrachat-200k
     output_dir: str = "/output/models"
-    hub_model_id: str = "devsecops/finetuned-llama"
+    hub_model_id: str = "devsecops/finetuned-llama-v2"
 
-    # LoRA
-    lora_r: int = 16
-    lora_alpha: int = 32
+    # LoRA (LoRA Without Regret optimal config)
+    lora_r: int = 256          # r=256 — sufficient capacity for SFT-scale datasets
+    lora_alpha: int = 16       # alpha=16 — stable scaling
     lora_dropout: float = 0.05
+    target_modules: str = "all-linear"  # ALL linear layers, not just attention
 
-    # Training
-    num_train_epochs: int = 3
-    per_device_train_batch_size: int = 4
-    gradient_accumulation_steps: int = 8  # effective batch = 32
-    learning_rate: float = 2e-4
+    # Training (LoRA Without Regret: batch < 32, LR = 2e-4)
+    num_train_epochs: int = 1  # 1 epoch sufficient for 940K dataset
+    per_device_train_batch_size: int = 2
+    gradient_accumulation_steps: int = 8  # effective batch = 16 (< 32!)
+    learning_rate: float = 2e-4  # 10x full FT rate
     max_seq_length: int = 2048
     warmup_ratio: float = 0.1
+    lr_scheduler_type: str = "cosine"
 
     # Optimization
     bf16: bool = True
     gradient_checkpointing: bool = True
     optim: str = "adamw_torch"
 
+    # Packing (LoRA Without Regret recommends packing=True)
+    packing: bool = True
+    packing_strategy: str = "bfd_split"  # Preserves all tokens
+
+    # Loss
+    assistant_only_loss: bool = True  # Only compute loss on assistant tokens
+
 
 def finetune(config: FinetuneConfig):
-    """Fine-tune a model with LoRA + SFT."""
+    """Fine-tune a model with LoRA + SFT (vNext — LoRA Without Regret config)."""
 
     # --- Trackio monitoring ---
     trackio.init(
         project="devsecops-ml",
-        name=f"sft-{config.model_name.split('/')[-1]}",
+        name=f"sft-{config.model_name.split('/')[-1]}-{config.dataset_key}",
         config=vars(config),
     )
 
-    # --- Quantization ---
-    bnb_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_quant_type="nf4",
-        bnb_4bit_compute_dtype=torch.bfloat16,
-        bnb_4bit_use_double_quant=True,
-    )
-
-    # --- Load model ---
-    tokenizer = AutoTokenizer.from_pretrained(
-        config.model_name,
-        trust_remote_code=True,
-        padding_side="right",
-    )
-    tokenizer.pad_token = tokenizer.eos_token
-
-    model = AutoModelForCausalLM.from_pretrained(
-        config.model_name,
-        quantization_config=bnb_config,
-        device_map="auto",
-        trust_remote_code=True,
-        torch_dtype=torch.bfloat16,
-    )
-    model = prepare_model_for_kbit_training(model)
+    # --- Dataset (best available) ---
+    dataset = load_and_prepare_dataset(config.dataset_key)
+    print(f"Dataset: {config.dataset_key} ({len(dataset)} examples)")
 
-    # --- LoRA ---
-    lora_config = LoraConfig(
+    # --- LoRA (LoRA Without Regret: all-linear, r=256) ---
+    peft_config = LoraConfig(
         r=config.lora_r,
         lora_alpha=config.lora_alpha,
         lora_dropout=config.lora_dropout,
         bias="none",
         task_type="CAUSAL_LM",
-        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
+        target_modules=config.target_modules,
     )
-    model = get_peft_model(model, lora_config)
-
-    # --- Dataset ---
-    dataset = load_dataset(config.dataset_name, split="train_sft[:5000]")
 
     # --- SFT Config ---
     sft_config = SFTConfig(
@@ -106,9 +151,14 @@ def finetune(config: FinetuneConfig):
         learning_rate=config.learning_rate,
         max_seq_length=config.max_seq_length,
         warmup_ratio=config.warmup_ratio,
+        lr_scheduler_type=config.lr_scheduler_type,
         bf16=config.bf16,
         gradient_checkpointing=config.gradient_checkpointing,
         optim=config.optim,
+        packing=config.packing,
+        packing_strategy=config.packing_strategy,
+        assistant_only_loss=config.assistant_only_loss,
+        eos_token="<|eot_id|>",
         logging_strategy="steps",
         logging_steps=10,
         logging_first_step=True,
@@ -121,12 +171,12 @@ def finetune(config: FinetuneConfig):
         disable_tqdm=True,
     )
 
-    # --- Trainer ---
+    # --- Trainer (SFTTrainer handles model loading + PEFT) ---
     trainer = SFTTrainer(
-        model=model,
-        args=sft_config,
+        model=config.model_name,
         train_dataset=dataset,
-        processing_class=tokenizer,
+        peft_config=peft_config,
+        args=sft_config,
     )
 
     # --- Train ---
@@ -135,7 +185,6 @@ def finetune(config: FinetuneConfig):
     # --- Save ---
     trainer.push_to_hub()
     trackio.finish()
-
     print(f"Model pushed to: https://huggingface.co/{config.hub_model_id}")