Upload scripts/train_qwen3_dpo_reasoning.py with huggingface_hub

scripts/train_qwen3_dpo_reasoning.py

@@ -0,0 +1,239 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "transformers>=4.45.0",
+#     "trl>=0.12.0",
+#     "peft>=0.13.0",
+#     "datasets>=3.0.0",
+#     "accelerate>=1.0.0",
+#     "huggingface_hub>=0.26.0",
+#     "torch>=2.4.0",
+#     "bitsandbytes>=0.44.0",
+# ]
+# [tool.uv]
+# index-strategy = "unsafe-best-match"
+# extra-index-url = ["https://download.pytorch.org/whl/cu124"]
+# ///
+"""
+DPO Training Script for Qwen3-0.6B on n8n Workflow Reasoning
+
+This script fine-tunes Qwen3-0.6B using Direct Preference Optimization (DPO)
+to improve reasoning quality when generating n8n workflows.
+
+The dataset contains:
+- prompt: task description for generating n8n workflow
+- chosen: high-quality response with detailed <thinking> reasoning
+- rejected: low-quality response with superficial reasoning or errors
+
+Usage:
+    hf jobs uv run \
+        --script train_qwen3_dpo_reasoning.py \
+        --flavor l40sx1 \
+        --name qwen3-dpo-reasoning \
+        --timeout 12h
+"""
+
+import os
+import torch
+from datasets import load_dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import LoraConfig
+from trl import DPOConfig, DPOTrainer
+from huggingface_hub import login
+
+# ============================================================================
+# CONFIGURATION
+# ============================================================================
+
+# Base model
+MODEL_NAME = os.environ.get("BASE_MODEL", "Qwen/Qwen3-0.6B")
+
+# Dataset
+DATASET_REPO = "stmasson/n8n-workflows-thinking"
+DATA_DIR = "data/dpo"
+
+# Output
+OUTPUT_DIR = "./qwen3-dpo-reasoning"
+HF_REPO = os.environ.get("HF_REPO", "stmasson/qwen3-0.6b-n8n-reasoning")
+
+# Hyperparameters
+NUM_EPOCHS = int(os.environ.get("NUM_EPOCHS", "1"))
+BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "1"))
+GRAD_ACCUM = int(os.environ.get("GRAD_ACCUM", "8"))
+LEARNING_RATE = float(os.environ.get("LEARNING_RATE", "5e-6"))
+MAX_LENGTH = int(os.environ.get("MAX_LENGTH", "4096"))
+MAX_PROMPT_LENGTH = int(os.environ.get("MAX_PROMPT_LENGTH", "512"))
+BETA = float(os.environ.get("BETA", "0.1"))  # DPO beta parameter
+
+# LoRA configuration
+LORA_R = int(os.environ.get("LORA_R", "32"))
+LORA_ALPHA = int(os.environ.get("LORA_ALPHA", "64"))
+LORA_DROPOUT = float(os.environ.get("LORA_DROPOUT", "0.05"))
+
+# ============================================================================
+# AUTHENTICATION
+# ============================================================================
+
+print("=" * 60)
+print("DPO TRAINING - QWEN3-0.6B N8N REASONING")
+print("=" * 60)
+
+hf_token = os.environ.get("HF_TOKEN")
+if hf_token:
+    login(token=hf_token)
+    print("Authenticated with HuggingFace")
+else:
+    print("Warning: HF_TOKEN not set, push disabled")
+
+# ============================================================================
+# LOAD MODEL AND TOKENIZER
+# ============================================================================
+
+print(f"\nLoading model: {MODEL_NAME}")
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.bfloat16,
+    attn_implementation="sdpa",
+    device_map="auto",
+    trust_remote_code=True,
+)
+
+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"  # Important for DPO
+
+print(f"Model loaded: {model.config.num_hidden_layers} layers, {model.config.hidden_size} hidden size")
+
+# ============================================================================
+# LORA CONFIGURATION
+# ============================================================================
+
+print(f"\nLoRA config: r={LORA_R}, alpha={LORA_ALPHA}")
+
+peft_config = LoraConfig(
+    r=LORA_R,
+    lora_alpha=LORA_ALPHA,
+    target_modules=[
+        "q_proj", "k_proj", "v_proj", "o_proj",
+        "gate_proj", "up_proj", "down_proj"
+    ],
+    lora_dropout=LORA_DROPOUT,
+    bias="none",
+    task_type="CAUSAL_LM"
+)
+
+# ============================================================================
+# LOAD DATASET
+# ============================================================================
+
+print(f"\nLoading dataset: {DATASET_REPO}")
+
+train_dataset = load_dataset(DATASET_REPO, data_dir=DATA_DIR, split="train")
+eval_dataset = load_dataset(DATASET_REPO, data_dir=DATA_DIR, split="validation")
+
+print(f"Train: {len(train_dataset)} examples")
+print(f"Validation: {len(eval_dataset)} examples")
+
+# Filter out extremely long examples to avoid OOM
+def filter_by_length(example):
+    prompt_len = len(example["prompt"])
+    chosen_len = len(example["chosen"])
+    rejected_len = len(example["rejected"])
+    # Filter examples where total chars > 50000 (roughly 12500 tokens)
+    return (prompt_len + max(chosen_len, rejected_len)) < 50000
+
+train_dataset = train_dataset.filter(filter_by_length)
+eval_dataset = eval_dataset.filter(filter_by_length)
+
+print(f"After filtering - Train: {len(train_dataset)}, Val: {len(eval_dataset)}")
+
+# Show example
+print("\nExample prompt:", train_dataset[0]["prompt"][:100], "...")
+
+# ============================================================================
+# DPO TRAINING CONFIGURATION
+# ============================================================================
+
+print(f"\nTraining configuration:")
+print(f"  - Epochs: {NUM_EPOCHS}")
+print(f"  - Batch size: {BATCH_SIZE}")
+print(f"  - Gradient accumulation: {GRAD_ACCUM}")
+print(f"  - Effective batch size: {BATCH_SIZE * GRAD_ACCUM}")
+print(f"  - Learning rate: {LEARNING_RATE}")
+print(f"  - Max length: {MAX_LENGTH}")
+print(f"  - DPO beta: {BETA}")
+
+training_args = DPOConfig(
+    output_dir=OUTPUT_DIR,
+    num_train_epochs=NUM_EPOCHS,
+    per_device_train_batch_size=BATCH_SIZE,
+    per_device_eval_batch_size=BATCH_SIZE,
+    gradient_accumulation_steps=GRAD_ACCUM,
+    learning_rate=LEARNING_RATE,
+    lr_scheduler_type="cosine",
+    warmup_ratio=0.1,
+    weight_decay=0.01,
+    bf16=True,
+    tf32=True,
+    logging_steps=10,
+    save_strategy="steps",
+    save_steps=500,
+    save_total_limit=3,
+    eval_strategy="steps",
+    eval_steps=500,
+    max_length=MAX_LENGTH,
+    max_prompt_length=MAX_PROMPT_LENGTH,
+    beta=BETA,
+    loss_type="sigmoid",  # Standard DPO loss
+    gradient_checkpointing=True,
+    gradient_checkpointing_kwargs={"use_reentrant": False},
+    report_to="none",
+    run_name="qwen3-dpo-reasoning",
+    hub_model_id=HF_REPO if hf_token else None,
+    push_to_hub=bool(hf_token),
+    hub_strategy="checkpoint",
+)
+
+# ============================================================================
+# TRAINING
+# ============================================================================
+
+print("\nInitializing DPO trainer...")
+
+trainer = DPOTrainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    peft_config=peft_config,
+    processing_class=tokenizer,
+)
+
+# Show trainable parameters
+trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+total_params = sum(p.numel() for p in model.parameters())
+print(f"\nTrainable parameters: {trainable_params:,} / {total_params:,} ({100 * trainable_params / total_params:.2f}%)")
+
+print("\n" + "=" * 60)
+print("STARTING DPO TRAINING")
+print("=" * 60)
+
+trainer.train()
+
+# ============================================================================
+# SAVE MODEL
+# ============================================================================
+
+print("\nSaving model...")
+trainer.save_model(f"{OUTPUT_DIR}/final")
+
+if hf_token:
+    print(f"Pushing to {HF_REPO}...")
+    trainer.push_to_hub()
+    print(f"Model available at: https://huggingface.co/{HF_REPO}")
+
+print("\n" + "=" * 60)
+print("DPO TRAINING COMPLETE")
+print("=" * 60)