File size: 4,739 Bytes
5748e4c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | """
Stage 1 SFT Training Script: Fine-tune Qwen2.5-Coder-7B-Instruct for
Python code generation + tool-calling + RAG-aware generation.
Based on:
- Qwen2.5-Coder Technical Report (arxiv:2409.12186): coarse-to-fine SFT recipe
- ToolACE (arxiv:2409.00920): LoRA r=16, alpha=32, LR=1e-4, cosine, 3 epochs
- Gorilla (arxiv:2305.15334): retriever-aware training pattern
- TRL v1.2.0 SFTTrainer with ChatML messages format
Reference implementation: TRL SFT docs (https://huggingface.co/docs/trl/sft_trainer)
Dataset format: conversational ChatML with "messages" column
Usage:
# Launch via hf_jobs (recommended)
# Or run locally:
python train_sft.py
"""
import os
import torch
import trackio
from datasets import load_dataset
from peft import LoraConfig
from trl import SFTConfig, SFTTrainer
# ============================================================================
# Configuration — EDIT THESE FOR YOUR USE CASE
# ============================================================================
BASE_MODEL = "Qwen/Qwen2.5-Coder-7B-Instruct"
DATASET_REPO = "your-username/code-toolcall-sft-data" # From prepare_data.py
OUTPUT_DIR = "./qwen25-coder-7b-code-toolcall"
HUB_MODEL_ID = "your-username/qwen25-coder-7b-code-toolcall"
HF_TOKEN = os.environ.get("HF_TOKEN")
# Training hyperparameters (from ToolACE + Qwen2.5-Coder papers)
LEARNING_RATE = 1e-4 # LoRA LR (10x FFT rate)
NUM_EPOCHS = 2 # 2 epochs (Magicoder recipe)
BATCH_SIZE = 2 # Per-device batch size
GRAD_ACCUM = 8 # Effective batch = 2 * 8 = 16
MAX_SEQ_LENGTH = 8192 # 8K context for tool-calling
WARMUP_RATIO = 0.1
LORA_R = 32 # Rank (32-64 for code tasks)
LORA_ALPHA = 64 # Alpha = 2 * r
LORA_DROPOUT = 0.05
# ============================================================================
# Trackio Monitoring
# ============================================================================
trackio.init(name="code-toolcall-sft", project="code-llm-finetuning")
# ============================================================================
# Load Dataset
# ============================================================================
print(f"Loading dataset: {DATASET_REPO}")
dataset = load_dataset(DATASET_REPO, split="train")
print(f"Dataset size: {len(dataset)} examples")
split = dataset.train_test_split(test_size=0.02, seed=42)
train_dataset = split["train"]
eval_dataset = split["test"]
print(f"Train: {len(train_dataset)}, Eval: {len(eval_dataset)}")
# ============================================================================
# LoRA Configuration
# ============================================================================
peft_config = LoraConfig(
r=LORA_R,
lora_alpha=LORA_ALPHA,
lora_dropout=LORA_DROPOUT,
bias="none",
task_type="CAUSAL_LM",
target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
"gate_proj", "up_proj", "down_proj"],
)
# ============================================================================
# SFT Training Configuration
# ============================================================================
training_args = SFTConfig(
output_dir=OUTPUT_DIR,
hub_model_id=HUB_MODEL_ID,
push_to_hub=True,
hub_token=HF_TOKEN,
num_train_epochs=NUM_EPOCHS,
per_device_train_batch_size=BATCH_SIZE,
gradient_accumulation_steps=GRAD_ACCUM,
learning_rate=LEARNING_RATE,
lr_scheduler_type="cosine",
warmup_ratio=WARMUP_RATIO,
weight_decay=0.01,
max_grad_norm=1.0,
bf16=True,
gradient_checkpointing=True,
max_seq_length=MAX_SEQ_LENGTH,
assistant_only_loss=True, # Only train on assistant responses
packing=False,
logging_strategy="steps",
logging_steps=10,
logging_first_step=True,
disable_tqdm=True,
eval_strategy="steps",
eval_steps=200,
per_device_eval_batch_size=2,
save_strategy="steps",
save_steps=500,
save_total_limit=3,
load_best_model_at_end=True,
seed=42,
report_to="none",
)
# ============================================================================
# Train
# ============================================================================
print(f"Loading model: {BASE_MODEL}")
trainer = SFTTrainer(
model=BASE_MODEL,
args=training_args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
peft_config=peft_config,
)
print("Starting training...")
result = trainer.train()
print(f"Training complete! Loss: {result.training_loss:.4f}")
trainer.push_to_hub(commit_message="Final model after SFT training")
print(f"Model pushed to: https://huggingface.co/{HUB_MODEL_ID}")
trackio.log({"final_loss": result.training_loss})
trackio.finish()
|