Add multi-GPU training script for 4xL40S

train_multi_gpu.py

@@ -0,0 +1,321 @@
+#!/usr/bin/env python3
+"""
+Qwen2.5-7B + glaive-function-calling-v2 QLoRA学習スクリプト
+マルチGPU対応版 (4xA10G等)
+
+実行方法:
+  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: バッチサイズを上げる
+    per_device_train_batch_size=4,  # 1GPUあたり4
+    per_device_eval_batch_size=4,
+    gradient_accumulation_steps=4,  # 有効バッチ: 4*4*num_gpus
+
+    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()