train.py

#!/usr/bin/env python3
"""
Qwen2.5-7B-Instruct + glaive-function-calling-v2 QLoRA学習スクリプト

目的: Function Calling能力の強化
データセット: glaiveai/glaive-function-calling-v2 (113k samples)
"""

import os
import sys
import time
from datetime import datetime

import torch
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer
from transformers.trainer_callback import TrainerCallback

# ============================================================
# 設定
# ============================================================
BASE_MODEL = "Qwen/Qwen2.5-7B-Instruct"
OUTPUT_MODEL_ID = "hajimemat/qwen2.5-7b-glaive-fc-lora-cloud"  # クラウド版
DATASET_NAME = "glaiveai/glaive-function-calling-v2"

# チェックポイント設定
CHECKPOINT_DIR = "./checkpoints"
FINAL_OUTPUT_DIR = "./output/final"

# ============================================================
# QLoRA量子化設定
# ============================================================
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
)

# ============================================================
# LoRA設定
# ============================================================
lora_config = LoraConfig(
    r=64,
    lora_alpha=16,
    lora_dropout=0.05,
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj"
    ],
    bias="none",
    task_type="CAUSAL_LM",
)


# ============================================================
# カスタムコールバック: 定期ログ出力
# ============================================================
class VerboseLoggingCallback(TrainerCallback):
    """詳細なログ出力用コールバック"""

    def __init__(self):
        self.start_time = None
        self.last_log_time = None

    def on_train_begin(self, args, state, control, **kwargs):
        self.start_time = time.time()
        self.last_log_time = self.start_time
        print("\n" + "=" * 70)
        print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Training started")
        print(f"  Total steps: {state.max_steps}")
        print(f"  Epochs: {args.num_train_epochs}")
        print(f"  Batch size: {args.per_device_train_batch_size} x {args.gradient_accumulation_steps}")
        print("=" * 70 + "\n")

    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs is None:
            return

        current_time = time.time()
        elapsed = current_time - self.start_time
        elapsed_str = time.strftime("%H:%M:%S", time.gmtime(elapsed))

        # 進捗計算
        progress = state.global_step / state.max_steps * 100 if state.max_steps > 0 else 0

        # ETA計算
        if state.global_step > 0:
            time_per_step = elapsed / state.global_step
            remaining_steps = state.max_steps - state.global_step
            eta_seconds = time_per_step * remaining_steps
            eta_str = time.strftime("%H:%M:%S", time.gmtime(eta_seconds))
        else:
            eta_str = "calculating..."

        # ログ出力
        loss = logs.get("loss", "N/A")
        lr = logs.get("learning_rate", "N/A")

        loss_str = f"{loss:.4f}" if isinstance(loss, float) else str(loss)
        lr_str = f"{lr:.2e}" if isinstance(lr, float) else str(lr)

        print(f"[{datetime.now().strftime('%H:%M:%S')}] "
              f"Step {state.global_step}/{state.max_steps} ({progress:.1f}%) | "
              f"Loss: {loss_str} | "
              f"LR: {lr_str} | "
              f"Elapsed: {elapsed_str} | ETA: {eta_str}")

        # GPU メモリ使用量（10ステップごと）
        if state.global_step % 100 == 0 and torch.cuda.is_available():
            allocated = torch.cuda.memory_allocated() / 1e9
            reserved = torch.cuda.memory_reserved() / 1e9
            print(f"         GPU Memory: {allocated:.2f}GB allocated, {reserved:.2f}GB reserved")

    def on_save(self, args, state, control, **kwargs):
        print(f"\n[{datetime.now().strftime('%H:%M:%S')}] "
              f"💾 Checkpoint saved at step {state.global_step}\n")

    def on_train_end(self, args, state, control, **kwargs):
        total_time = time.time() - self.start_time
        total_str = time.strftime("%H:%M:%S", time.gmtime(total_time))
        print("\n" + "=" * 70)
        print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Training completed!")
        print(f"  Total time: {total_str}")
        print(f"  Final step: {state.global_step}")
        print("=" * 70 + "\n")


# ============================================================
# データセット変換
# ============================================================
def convert_glaive_to_chatml(example: dict) -> dict:
    """
    glaive-function-calling-v2形式をChatML形式に変換

    元データ形式:
    - system: 関数定義を含むシステムプロンプト
    - chat: "USER: ... ASSISTANT: ..." 形式の会話
    """
    parts = []

    # システムプロンプト
    if example.get("system"):
        parts.append(f"<|im_start|>system\n{example['system']}<|im_end|>")

    # 会話を解析
    chat = example.get("chat", "")
    if chat:
        # "USER:" と "ASSISTANT:" で分割
        # 複数ターンに対応
        current_role = None
        current_content = []

        for line in chat.split("\n"):
            line = line.strip()
            if line.startswith("USER:"):
                # 前のメッセージを保存
                if current_role and current_content:
                    content = "\n".join(current_content).strip()
                    if content:
                        parts.append(f"<|im_start|>{current_role}\n{content}<|im_end|>")
                current_role = "user"
                current_content = [line[5:].strip()]  # "USER:" を除去
            elif line.startswith("ASSISTANT:"):
                # 前のメッセージを保存
                if current_role and current_content:
                    content = "\n".join(current_content).strip()
                    if content:
                        parts.append(f"<|im_start|>{current_role}\n{content}<|im_end|>")
                current_role = "assistant"
                current_content = [line[10:].strip()]  # "ASSISTANT:" を除去
            elif current_role:
                current_content.append(line)

        # 最後のメッセージを保存
        if current_role and current_content:
            content = "\n".join(current_content).strip()
            if content:
                parts.append(f"<|im_start|>{current_role}\n{content}<|im_end|>")

    return {"text": "\n".join(parts)}


def load_and_prepare_dataset():
    """データセットを読み込んで前処理"""
    print(f"\n{'=' * 60}")
    print(f"Loading dataset: {DATASET_NAME}")
    print(f"{'=' * 60}")

    # データセット読み込み
    dataset = load_dataset(DATASET_NAME, split="train")
    print(f"Original size: {len(dataset)} examples")

    # 変換
    print("Converting to ChatML format...")
    dataset = dataset.map(
        convert_glaive_to_chatml,
        remove_columns=dataset.column_names,
        num_proc=4,
        desc="Converting"
    )

    # 空のサンプルをフィルタ
    dataset = dataset.filter(lambda x: len(x["text"]) > 50)
    print(f"After filtering: {len(dataset)} examples")

    # サンプル表示
    print("\n--- Sample data ---")
    sample = dataset[0]["text"]
    print(sample[:500] + "..." if len(sample) > 500 else sample)
    print("--- End sample ---\n")

    # シャッフルしてTrain/Test分割
    dataset = dataset.shuffle(seed=42)
    split = dataset.train_test_split(test_size=0.02, seed=42)

    print(f"Train: {len(split['train'])} examples")
    print(f"Test: {len(split['test'])} examples")

    return split


# ============================================================
# 学習パラメータ
# ============================================================
training_args = TrainingArguments(
    output_dir=CHECKPOINT_DIR,

    # エポック・ステップ
    num_train_epochs=1,
    max_steps=-1,  # -1 = エポックベース

    # バッチサイズ (L40S 48GB - 攻めた設定)
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    gradient_accumulation_steps=2,  # 有効バッチサイズ: 8*2=16

    # 学習率 (1エポックで収束するよう高め)
    learning_rate=2e-4,
    weight_decay=0.01,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",

    # 最適化
    optim="paged_adamw_8bit",
    fp16=False,
    bf16=True,
    max_grad_norm=0.3,

    # ログ・保存（重要！）
    logging_steps=10,           # 10ステップごとにログ
    save_steps=500,             # 500ステップごとにチェックポイント
    save_total_limit=3,         # 最新3つのチェックポイントを保持
    eval_strategy="steps",
    eval_steps=500,             # 500ステップごとに評価

    # その他
    report_to="none",
    group_by_length=True,
    gradient_checkpointing=False,  # L40Sは48GBあるのでオフで高速化
    torch_compile=False,  # 初回コンパイル時間を避ける
    dataloader_num_workers=4,
    dataloader_pin_memory=True,

    # 再開用
    save_safetensors=True,
    load_best_model_at_end=False,
)


# ============================================================
# メイン
# ============================================================
def main():
    print("\n" + "=" * 70)
    print("  Qwen2.5-7B + glaive-function-calling-v2 QLoRA Training")
    print("=" * 70)
    print(f"Start time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"Base model: {BASE_MODEL}")
    print(f"Dataset: {DATASET_NAME}")
    print(f"Output: {OUTPUT_MODEL_ID}")
    print("=" * 70 + "\n")

    # GPU確認
    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name(0)
        gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9
        print(f"GPU: {gpu_name}")
        print(f"VRAM: {gpu_mem:.1f} GB")
    else:
        print("ERROR: No GPU available!")
        sys.exit(1)

    # データセット読み込み
    dataset = load_and_prepare_dataset()

    # トークナイザー読み込み
    print(f"\nLoading tokenizer: {BASE_MODEL}")
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True)
    tokenizer.padding_side = "right"
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # モデル読み込み (4bit量子化)
    print(f"\nLoading model: {BASE_MODEL} (4-bit quantized)")
    print("This may take a few minutes...")
    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        quantization_config=bnb_config,
        device_map="auto",
        attn_implementation="sdpa",
        trust_remote_code=True,
    )

    # 学習準備
    print("\nPreparing model for training...")
    model = prepare_model_for_kbit_training(model)
    model = get_peft_model(model, lora_config)

    print("\nTrainable parameters:")
    model.print_trainable_parameters()

    # SFTTrainer設定
    trainer = SFTTrainer(
        model=model,
        train_dataset=dataset["train"],
        eval_dataset=dataset["test"],
        args=training_args,
        peft_config=lora_config,
        tokenizer=tokenizer,
        max_seq_length=1024,  # 高速化のため短縮
        packing=False,  # flash attention不要で高速化
        dataset_text_field="text",
        callbacks=[VerboseLoggingCallback()],
    )

    # チェックポイントからの再開確認
    resume_from = None
    if os.path.exists(CHECKPOINT_DIR):
        checkpoints = [d for d in os.listdir(CHECKPOINT_DIR) if d.startswith("checkpoint-")]
        if checkpoints:
            latest = max(checkpoints, key=lambda x: int(x.split("-")[1]))
            resume_from = os.path.join(CHECKPOINT_DIR, latest)
            print(f"\n📂 Found checkpoint: {resume_from}")
            print("   Resuming from checkpoint...")

    # 学習実行
    print("\n" + "=" * 70)
    print("Starting training...")
    print("=" * 70)

    trainer.train(resume_from_checkpoint=resume_from)

    # 最終モデル保存
    print(f"\nSaving final model to {FINAL_OUTPUT_DIR}...")
    trainer.save_model(FINAL_OUTPUT_DIR)
    tokenizer.save_pretrained(FINAL_OUTPUT_DIR)

    # HFにアップロード
    print(f"\nUploading to HuggingFace: {OUTPUT_MODEL_ID}")
    try:
        trainer.model.push_to_hub(OUTPUT_MODEL_ID, private=True)
        tokenizer.push_to_hub(OUTPUT_MODEL_ID, private=True)
        print(f"✅ Model uploaded to: https://huggingface.co/{OUTPUT_MODEL_ID}")
    except Exception as e:
        print(f"⚠️ Upload failed: {e}")
        print("   Model saved locally. Please upload manually.")

    print("\n" + "=" * 70)
    print("🎉 Training complete!")
    print("=" * 70)


if __name__ == "__main__":
    main()