train_multi_gpu.py

#!/usr/bin/env python3
"""
Qwen2.5-7B + glaive-function-calling-v2 QLoRA学習スクリプト
マルチGPU対応版 (4xL40S等)

実行方法:
  accelerate launch --num_processes 4 train_multi_gpu.py
"""

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"
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

    def on_train_begin(self, args, state, control, **kwargs):
        self.start_time = time.time()
        if state.is_world_process_zero:
            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"  Num GPUs: {args.world_size}")
            print(f"  Per device batch: {args.per_device_train_batch_size}")
            print(f"  Gradient accum: {args.gradient_accumulation_steps}")
            print(f"  Effective batch: {args.per_device_train_batch_size * args.gradient_accumulation_steps * args.world_size}")
            print("=" * 70 + "\n")

    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs is None or not state.is_world_process_zero:
            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

        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")

        print(f"[{datetime.now().strftime('%H:%M:%S')}] "
              f"Step {state.global_step}/{state.max_steps} ({progress:.1f}%) | "
              f"Loss: {loss:.4f if isinstance(loss, float) else loss} | "
              f"LR: {lr:.2e if isinstance(lr, float) else lr} | "
              f"Elapsed: {elapsed_str} | ETA: {eta_str}")

    def on_save(self, args, state, control, **kwargs):
        if state.is_world_process_zero:
            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):
        if state.is_world_process_zero:
            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("=" * 70 + "\n")


# ============================================================
# データセット変換
# ============================================================
def convert_glaive_to_chatml(example: dict) -> dict:
    parts = []

    if example.get("system"):
        parts.append(f"<|im_start|>system\n{example['system']}<|im_end|>")

    chat = example.get("chat", "")
    if chat:
        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()]
            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()]
            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"\nLoading dataset: {DATASET_NAME}")

    dataset = load_dataset(DATASET_NAME, split="train")
    print(f"Original size: {len(dataset)} examples")

    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")

    dataset = dataset.shuffle(seed=42)
    split = dataset.train_test_split(test_size=0.02, seed=42)

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


# ============================================================
# 学習パラメータ（マルチGPU最適化）
# ============================================================
num_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 1

training_args = TrainingArguments(
    output_dir=CHECKPOINT_DIR,

    num_train_epochs=2,

    # マルチGPU: L40Sは48GB VRAMなのでバッチサイズを上げる
    per_device_train_batch_size=8,  # 1GPUあたり8 (L40S 48GB)
    per_device_eval_batch_size=8,
    gradient_accumulation_steps=2,  # 有効バッチ: 8*2*4=64

    learning_rate=1e-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,
    save_steps=500,
    save_total_limit=3,
    eval_strategy="steps",
    eval_steps=500,

    report_to="none",
    group_by_length=True,
    gradient_checkpointing=True,

    # マルチGPU設定
    ddp_find_unused_parameters=False,
    dataloader_num_workers=4,

    save_safetensors=True,
)


# ============================================================
# メイン
# ============================================================
def main():
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    is_main = local_rank == 0

    if is_main:
        print("\n" + "=" * 70)
        print("  Qwen2.5-7B + glaive-function-calling-v2 QLoRA Training")
        print("  Multi-GPU Version")
        print("=" * 70)
        print(f"Start: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print(f"GPUs available: {torch.cuda.device_count()}")
        for i in range(torch.cuda.device_count()):
            print(f"  GPU {i}: {torch.cuda.get_device_name(i)}")
        print("=" * 70 + "\n")

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

    # トークナイザー
    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

    # モデル
    if is_main:
        print(f"\nLoading model: {BASE_MODEL}")

    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        quantization_config=bnb_config,
        device_map={"": local_rank},  # 各GPUに配置
        attn_implementation="sdpa",
        trust_remote_code=True,
    )

    model = prepare_model_for_kbit_training(model)
    model = get_peft_model(model, lora_config)

    if is_main:
        model.print_trainable_parameters()

    # Trainer
    trainer = SFTTrainer(
        model=model,
        train_dataset=dataset["train"],
        eval_dataset=dataset["test"],
        args=training_args,
        peft_config=lora_config,
        processing_class=tokenizer,
        max_seq_length=2048,
        packing=True,
        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)
            if is_main:
                print(f"\n📂 Resuming from: {resume_from}")

    # 学習
    trainer.train(resume_from_checkpoint=resume_from)

    # 保存（メインプロセスのみ）
    if is_main:
        print(f"\nSaving to {FINAL_OUTPUT_DIR}...")
        trainer.save_model(FINAL_OUTPUT_DIR)
        tokenizer.save_pretrained(FINAL_OUTPUT_DIR)

        print(f"\nUploading to: {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"✅ Uploaded: https://huggingface.co/{OUTPUT_MODEL_ID}")
        except Exception as e:
            print(f"⚠️ Upload failed: {e}")

        print("\n🎉 Training complete!")


if __name__ == "__main__":
    main()