code_llm_train.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Code LLM - QLoRA Fine-tuning Script
====================================
Base Model: Qwen/Qwen2.5-Coder-3B
Method: QLoRA SFT (4-bit NF4 + LoRA r=64)
Datasets: Code-Feedback (66K) + Magicoder-OSS (75K) + Evol-CodeAlpaca (110K) = ~250K

Hardware: RTX 3070 (8GB VRAM) or any GPU >= 8GB
Training time: ~6-8 hours (3 epochs)

Usage:
    pip install -r requirements_code.txt
    python code_llm_train.py
"""

import os
import sys
import torch
from datetime import datetime

# ============================================================
#  CONFIGURATION - 請修改這裡
# ============================================================
MODEL_NAME = "Qwen/Qwen2.5-Coder-3B"
HF_USERNAME = "YOUR_HF_USERNAME"       # 改成你的 HuggingFace 用戶名

# 訓練超參數 (RTX 3070 8GB 優化)
TRAINING_CONFIG = {
    "learning_rate": 2e-4,
    "num_epochs": 3,
    "batch_size": 1,
    "gradient_accumulation": 16,
    "max_seq_length": 2048,
    "lora_r": 64,
    "lora_alpha": 128,
    "lora_dropout": 0.05,
    "warmup_ratio": 0.05,
}

OUTPUT_DIR = f"{HF_USERNAME}/code-qwen2.5-coder-3b"
# ============================================================


def print_banner(text):
    print(f"\n{'='*60}")
    print(f"  {text}")
    print(f"{'='*60}")


def check_environment():
    print_banner("ENVIRONMENT CHECK")
    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name(0)
        vram = torch.cuda.get_device_properties(0).total_mem / 1024**3
        print(f"✅ GPU: {gpu_name}")
        print(f"   VRAM: {vram:.1f} GB")
        if vram < 7:
            print("⚠️  VRAM < 8GB, 可能會 OOM，建議降低 max_seq_length 到 1024")
    else:
        print("❌ 沒有偵測到 GPU！此腳本需要 NVIDIA GPU")
        sys.exit(1)

    required = ["transformers", "trl", "peft", "bitsandbytes", "accelerate", "datasets"]
    missing = []
    for pkg in required:
        try:
            __import__(pkg)
            print(f"✅ {pkg}")
        except ImportError:
            missing.append(pkg)
            print(f"❌ {pkg}")
    if missing:
        print(f"\n請運行: pip install {' '.join(missing)}")
        sys.exit(1)


def load_datasets():
    from datasets import load_dataset, concatenate_datasets

    print_banner("LOADING DATASETS")

    print("📦 [1/3] Code-Feedback (66K multi-turn coding chat)...")
    code_feedback = load_dataset("m-a-p/Code-Feedback", split="train")
    cf_msgs = code_feedback.map(
        lambda x: {"messages": x["messages"]},
        remove_columns=[c for c in code_feedback.column_names if c != "messages"],
    )
    print(f"   ✅ {len(cf_msgs)} samples loaded")

    print("📦 [2/3] Magicoder-OSS-Instruct (75K real GitHub seeds)...")
    magicoder = load_dataset("ise-uiuc/Magicoder-OSS-Instruct-75K", split="train")
    def convert_magicoder(example):
        return {"messages": [
            {"role": "system", "content": "You are an exceptionally skilled programmer. Write clean, efficient, well-documented code."},
            {"role": "user", "content": example["problem"]},
            {"role": "assistant", "content": example["solution"]},
        ]}
    mc_msgs = magicoder.map(convert_magicoder, remove_columns=magicoder.column_names)
    print(f"   ✅ {len(mc_msgs)} samples converted")

    print("📦 [3/3] Evol-CodeAlpaca (110K complexity-evolved)...")
    evol = load_dataset("theblackcat102/evol-codealpaca-v1", split="train")
    def convert_evol(example):
        return {"messages": [
            {"role": "system", "content": "You are an exceptionally skilled programmer. Write clean, efficient, well-documented code."},
            {"role": "user", "content": example["instruction"]},
            {"role": "assistant", "content": example["output"]},
        ]}
    evol_msgs = evol.map(convert_evol, remove_columns=evol.column_names)
    print(f"   ✅ {len(evol_msgs)} samples converted")

    print("\n🔄 合併數據集...")
    combined = concatenate_datasets([cf_msgs, mc_msgs, evol_msgs]).shuffle(seed=42)
    split = combined.train_test_split(test_size=0.02, seed=42)
    train_ds, eval_ds = split["train"], split["test"]

    print(f"\n📊 數據集統計:")
    print(f"   Code-Feedback:  {len(cf_msgs):>7,} samples")
    print(f"   Magicoder-OSS:  {len(mc_msgs):>7,} samples")
    print(f"   Evol-CodeAlpaca:{len(evol_msgs):>7,} samples")
    print(f"   {'─'*35}")
    print(f"   總計訓練:       {len(train_ds):>7,} samples")
    print(f"   總計驗證:       {len(eval_ds):>7,} samples")
    return train_ds, eval_ds


def setup_model():
    from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
    from peft import LoraConfig, prepare_model_for_kbit_training, get_peft_model

    print_banner("LOADING MODEL")
    print(f"🤖 Model: {MODEL_NAME}")

    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"
    print(f"   Vocab: {len(tokenizer):,} tokens")

    print("\n⚡ 配置 QLoRA (4-bit NF4 + double quant)...")
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True, bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True,
    )

    print("📥 載入模型...")
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME, quantization_config=bnb_config, device_map="auto", trust_remote_code=True,
    )
    model = prepare_model_for_kbit_training(model)
    print("✅ 模型準備完成")

    print(f"\n🔧 配置 LoRA (r={TRAINING_CONFIG['lora_r']}, alpha={TRAINING_CONFIG['lora_alpha']})...")
    lora_config = LoraConfig(
        r=TRAINING_CONFIG["lora_r"], lora_alpha=TRAINING_CONFIG["lora_alpha"],
        lora_dropout=TRAINING_CONFIG["lora_dropout"], bias="none", task_type="CAUSAL_LM",
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        modules_to_save=["lm_head", "embed_tokens"],
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()
    return model, tokenizer, lora_config


def create_trainer(model, tokenizer, train_ds, eval_ds, lora_config):
    from trl import SFTTrainer, SFTConfig

    print_banner("CONFIGURING TRAINER")
    run_name = f"code-qwen-{datetime.now().strftime('%m%d-%H%M')}"

    report_to = []
    try:
        import trackio
        trackio.init(project="code-llm", experiment="qlora-sft", run_name=run_name)
        report_to = ["trackio"]
        print("✅ Trackio 監控已啟動")
    except Exception:
        print("⚠️  Trackio 不可用，使用 tensorboard")
        report_to = ["tensorboard"]

    training_args = SFTConfig(
        learning_rate=TRAINING_CONFIG["learning_rate"], lr_scheduler_type="cosine",
        warmup_ratio=TRAINING_CONFIG["warmup_ratio"],
        num_train_epochs=TRAINING_CONFIG["num_epochs"],
        per_device_train_batch_size=TRAINING_CONFIG["batch_size"],
        gradient_accumulation_steps=TRAINING_CONFIG["gradient_accumulation"],
        max_seq_length=TRAINING_CONFIG["max_seq_length"],
        gradient_checkpointing=True, bf16=True, fp16=False,
        optim="paged_adamw_8bit", packing=True,
        output_dir="./output_code", logging_steps=10, save_steps=1000, save_total_limit=2,
        eval_strategy="steps", eval_steps=1000,
        push_to_hub=True, hub_model_id=OUTPUT_DIR, hub_strategy="checkpoint",
        report_to=report_to, logging_strategy="steps", logging_first_step=True,
        remove_unused_columns=False, dataloader_num_workers=4, seed=42,
    )

    trainer = SFTTrainer(
        model=model, args=training_args, train_dataset=train_ds, eval_dataset=eval_ds,
        processing_class=tokenizer, peft_config=lora_config,
    )

    total_steps = len(train_ds) // (TRAINING_CONFIG["batch_size"] * TRAINING_CONFIG["gradient_accumulation"]) * TRAINING_CONFIG["num_epochs"]
    print(f"\n📋 訓練計劃:")
    print(f"   數據量:     {len(train_ds):,} samples")
    print(f"   Batch size: {TRAINING_CONFIG['batch_size']} × {TRAINING_CONFIG['gradient_accumulation']} = {TRAINING_CONFIG['batch_size'] * TRAINING_CONFIG['gradient_accumulation']}")
    print(f"   Epochs:     {TRAINING_CONFIG['num_epochs']}")
    print(f"   預估步數:   ~{total_steps:,} steps")
    print(f"   Packing:    ✅ 啟用")
    print(f"   Optimizer:  paged_adamw_8bit")
    print(f"   輸出位置:   https://huggingface.co/{OUTPUT_DIR}")
    return trainer, run_name


def train(trainer):
    print_banner("TRAINING")
    print("🚀 開始訓練...\n   按 Ctrl+C 可隨時中斷並保存\n")
    try:
        result = trainer.train()
        print(f"\n✅ 訓練完成！ Steps: {result.global_step}, Loss: {result.training_loss:.4f}")
        return True
    except KeyboardInterrupt:
        print("\n⚠️  訓練被中斷，正在保存...")
        trainer.save_model()
        return True
    except Exception as e:
        print(f"\n❌ 訓練失敗: {e}")
        raise


def save_and_push(trainer):
    print_banner("SAVING & UPLOADING")
    try:
        print("📤 上傳模型到 HuggingFace Hub...")
        trainer.push_to_hub()
        print(f"\n✅ 模型已上傳!\n🔗 https://huggingface.co/{OUTPUT_DIR}")
    except Exception as e:
        print(f"⚠️  上傳失敗: {e}\n   模型已保存在 ./output_code 目錄")


def main():
    print("""
    ╔════════════════════════════════════════════════════════════╗
    ║          Code LLM - QLoRA Fine-tuning                     ║
    ║          Base: Qwen2.5-Coder-3B                           ║
    ║          Data: 250K code samples (3 datasets)             ║
    ╚════════════════════════════════════════════════════════════╝
    """)
    check_environment()
    train_ds, eval_ds = load_datasets()
    model, tokenizer, lora_config = setup_model()
    trainer, run_name = create_trainer(model, tokenizer, train_ds, eval_ds, lora_config)
    success = train(trainer)
    if success:
        save_and_push(trainer)
    print_banner("DONE")
    print(f"  Run: {run_name}\n  Model: https://huggingface.co/{OUTPUT_DIR}")

if __name__ == "__main__":
    main()