예제 소스

train_vlm.py

#!/usr/bin/env python3
"""Qwen3-VL LoRA SFT 학습 — `dataset_omook.py` 가 만든 train.jsonl 사용.

예)
  python train_vlm.py \\
      --dataset-jsonl sft_data/omook_sft/train.jsonl \\
      --base-model /Users/.../vlm_models/Qwen__Qwen3-VL-8B-Instruct \\
      --output-dir vlm_models/lora_omook \\
      --epochs 1 --batch-size 1 --grad-accum 8

학습 후, 서버에서 다음과 같이 LoRA 어댑터를 얹어 추론할 수 있습니다.
  python vlm_pva_server.py --model-id <base> --lora-dir vlm_models/lora_omook ...

주의 (Apple Silicon):
  - 8B 모델은 MPS 에서도 메모리/속도 제약이 큽니다. 가능하면 Qwen3-VL-2B 로 시작하세요.
  - bf16 은 CUDA 만, MPS 에서는 fp16 으로 자동 폴백합니다.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from typing import Any, Dict, List

# 이 스크립트가 사는 디렉토리(omook/) 안에 'datasets/' 라는 폴더가 있으면 Python 이
# 그것을 namespace package 로 잡아 huggingface `datasets` 패키지를 가린다. 그 결과
# `from datasets import Dataset` 가 ImportError 로 죽는다. import 전에 sys.path 에서
# 스크립트 디렉토리/cwd 를 빼서 PyPI 패키지 우선으로 해결한다.
_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path[:] = [p for p in sys.path if os.path.abspath(p or ".") != _HERE]

from PIL import Image  # noqa: E402


def _load_examples(path: str) -> List[Dict[str, Any]]:
    out: List[Dict[str, Any]] = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            rec = json.loads(line)
            if "image_path" in rec and not os.path.exists(rec["image_path"]):
                # 이미지 파일이 없으면 학습 불가능 → 스킵 (또는 dataset_omook 재실행).
                continue
            out.append(rec)
    return out


def main() -> None:
    parser = argparse.ArgumentParser(description="LoRA SFT for omook VLM")
    parser.add_argument("--dataset-jsonl", required=True)
    parser.add_argument("--base-model", required=True, help="HF id 또는 로컬 경로")
    parser.add_argument("--output-dir", default="vlm_models/lora_omook")
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--batch-size", type=int, default=1)
    parser.add_argument("--grad-accum", type=int, default=8)
    parser.add_argument(
        "--lr",
        type=float,
        default=5e-5,
        help="LoRA SFT LR. 8B+에서 2e-4 는 collapse 위험이 큼 → 5e-5 권장.",
    )
    parser.add_argument(
        "--warmup-ratio",
        type=float,
        default=0.05,
        help="총 step 중 워밍업 비율 (0.05 = 5퍼센트). 0 = warmup 없음.",
    )
    parser.add_argument(
        "--max-grad-norm",
        type=float,
        default=1.0,
        help="gradient clipping. 발산 방지.",
    )
    parser.add_argument("--lora-r", type=int, default=8)
    parser.add_argument("--lora-alpha", type=int, default=16)
    parser.add_argument("--lora-dropout", type=float, default=0.05)
    parser.add_argument("--max-samples", type=int, default=0, help="0 = 전체")
    parser.add_argument("--logging-steps", type=int, default=10)
    parser.add_argument("--save-steps", type=int, default=200)
    parser.add_argument(
        "--num-workers",
        type=int,
        default=max(2, (os.cpu_count() or 4) // 2),
        help="DataLoader 워커 수 (기본: cpu/2). MPS 에서는 fork 호환성 때문에 0~2 권장.",
    )
    parser.add_argument(
        "--torch-threads",
        type=int,
        default=0,
        help="0 = torch 기본. 양수면 torch.set_num_threads(N) 로 CPU 텐서 op 병렬도 지정.",
    )
    parser.add_argument(
        "--device",
        default="auto",
        choices=("auto", "cuda", "mps", "cpu"),
        help="auto = 가능한 가속기 자동 선택",
    )
    parser.add_argument(
        "--target-modules",
        nargs="+",
        default=["q_proj", "k_proj", "v_proj", "o_proj"],
        help="LoRA 적용 모듈 이름 (Qwen 계열 LM 어텐션 기본값)",
    )
    parser.add_argument(
        "--max-seq-len",
        type=int,
        default=0,
        help="0 = truncation 끔(권장: Qwen3-VL 은 잘리면 이미지 토큰 수≠비전 패치로 터짐). "
        "양수면 max_length 로 자름(위험).",
    )
    args = parser.parse_args()

    import torch
    from datasets import Dataset
    from peft import LoraConfig, get_peft_model
    from transformers import (
        AutoModelForImageTextToText,
        AutoProcessor,
        Trainer,
        TrainingArguments,
    )

    if args.device == "auto":
        if torch.cuda.is_available():
            device = "cuda"
        elif torch.backends.mps.is_available():
            device = "mps"
        else:
            device = "cpu"
    else:
        device = args.device

    if args.torch_threads > 0:
        torch.set_num_threads(args.torch_threads)
    print(
        f"[train] device={device} num_workers={args.num_workers} "
        f"torch_threads={torch.get_num_threads()}"
    )

    examples = _load_examples(args.dataset_jsonl)
    if args.max_samples and len(examples) > args.max_samples:
        examples = examples[: args.max_samples]
    if not examples:
        raise SystemExit("학습 예제가 0건입니다. dataset_omook.py 출력과 image_path 를 확인하세요.")
    print(f"[train] examples={len(examples)}")

    ds = Dataset.from_list(examples)

    print(f"[train] loading processor: {args.base_model}")
    processor = AutoProcessor.from_pretrained(args.base_model, trust_remote_code=True)

    if device == "cuda":
        dtype = torch.bfloat16
    elif device == "mps":
        dtype = torch.float16
    else:
        dtype = torch.float32

    print(f"[train] loading model dtype={dtype}")
    model = AutoModelForImageTextToText.from_pretrained(
        args.base_model,
        dtype=dtype,
        trust_remote_code=True,
    )
    if device != "cuda":
        # cuda 는 device_map=auto 도 가능하지만 단순화를 위해 명시 이동.
        model.to(device)
    # gradient_checkpointing 과 use_cache 는 동시 사용 불가 — 미리 꺼서 경고를 없앤다.
    if hasattr(model, "config"):
        model.config.use_cache = False

    lora_cfg = LoraConfig(
        r=args.lora_r,
        lora_alpha=args.lora_alpha,
        target_modules=list(args.target_modules),
        lora_dropout=args.lora_dropout,
        bias="none",
        task_type="CAUSAL_LM",
    )
    model = get_peft_model(model, lora_cfg)
    model.print_trainable_parameters()

    pad_id = processor.tokenizer.pad_token_id
    if pad_id is None:
        # Qwen 토크나이저는 보통 eos 와 동일한 pad 를 가짐. 없으면 명시적으로 지정.
        pad_id = processor.tokenizer.eos_token_id
        processor.tokenizer.pad_token_id = pad_id

    def collate(batch: List[Dict[str, Any]]):
        """SFT 콜레이트 — *어시스턴트 응답 구간만* loss 에 포함.

        loss 마스킹을 안 하면 모델이 긴 user prompt(보드 텍스트 등) 전체를 외우려 들다가
        단일 토큰만 출력하는 collapse(예: "!!!!!") 로 빠진다. user 측만 적용한
        chat template 의 토큰 길이를 측정해 그 앞쪽을 -100 으로 마스킹한다.
        """
        texts_full: List[str] = []
        texts_user_only: List[str] = []
        images: List[List[Image.Image]] = []
        for rec in batch:
            img = None
            if "image_path" in rec:
                img = Image.open(rec["image_path"]).convert("RGB")
            elif "image" in rec:
                img = rec["image"]
            if img is None:
                raise RuntimeError(f"이미지 없음: id={rec.get('id')}")
            user_msg = {
                "role": "user",
                "content": [
                    {"type": "image", "image": img},
                    {"type": "text", "text": rec["prompt"]},
                ],
            }
            messages_full = [
                user_msg,
                {"role": "assistant", "content": rec["answer"]},
            ]
            messages_user = [user_msg]
            full = processor.apply_chat_template(
                messages_full, tokenize=False, add_generation_prompt=False
            )
            user_only = processor.apply_chat_template(
                messages_user, tokenize=False, add_generation_prompt=True
            )
            texts_full.append(full)
            texts_user_only.append(user_only)
            images.append([img])

        if args.max_seq_len and args.max_seq_len > 0:
            trunc_kw: Dict[str, Any] = {
                "truncation": True,
                "max_length": args.max_seq_len,
            }
        else:
            trunc_kw = {"truncation": False}

        out = processor(
            text=texts_full,
            images=images,
            padding=True,
            return_tensors="pt",
            **trunc_kw,
        )

        # 동일한 이미지 컨텍스트로 user-only 길이를 측정 (이미지 토큰 수가 같아야 prefix
        # 길이가 일치한다).
        prefix_meas = processor(
            text=texts_user_only,
            images=images,
            padding=False,
            return_tensors=None,
            **trunc_kw,
        )
        prefix_lens = [len(ids) for ids in prefix_meas["input_ids"]]

        labels = out["input_ids"].clone()
        if pad_id is not None:
            labels[labels == pad_id] = -100
        # user prompt 구간 마스킹.
        for i, plen in enumerate(prefix_lens):
            plen = min(plen, labels.shape[1])
            labels[i, :plen] = -100
        out["labels"] = labels
        return out

    # MPS 는 pinned memory 미지원 → 명시적으로 꺼서 경고를 없앤다.
    dataloader_pin_memory = (device == "cuda")

    # warmup_ratio 는 transformers v5.2 에서 제거 예정 → 총 step 기반 warmup_steps 로 변환.
    steps_per_epoch = max(1, (len(examples) + args.batch_size - 1) // args.batch_size)
    total_steps = max(1, (steps_per_epoch * args.epochs) // max(1, args.grad_accum))
    warmup_steps = max(0, int(round(total_steps * args.warmup_ratio)))

    targs = TrainingArguments(
        output_dir=args.output_dir,
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.batch_size,
        gradient_accumulation_steps=args.grad_accum,
        learning_rate=args.lr,
        warmup_steps=warmup_steps,
        max_grad_norm=args.max_grad_norm,
        logging_steps=args.logging_steps,
        save_steps=args.save_steps,
        save_total_limit=2,
        bf16=(dtype == torch.bfloat16),
        fp16=(dtype == torch.float16),
        remove_unused_columns=False,
        report_to=[],
        dataloader_num_workers=args.num_workers,
        dataloader_persistent_workers=(args.num_workers > 0),
        dataloader_pin_memory=dataloader_pin_memory,
        gradient_checkpointing=True,
    )

    trainer = Trainer(
        model=model,
        args=targs,
        train_dataset=ds,
        data_collator=collate,
    )
    trainer.train()

    # ---- 발산/NaN 게이트 ----
    # log_history 의 모든 step 손실/그래드를 훑어서 NaN/Inf/극단값이 보이면 LoRA 저장을 막는다.
    # collapsed LoRA 가 디스크에 들어가는 것을 차단해서 추론 모델이 다음 라운드까지 멍청해지지 않게 한다.
    import math

    bad_reason = None
    last_loss = None
    last_gn = None
    for entry in trainer.state.log_history:
        loss = entry.get("loss")
        gn = entry.get("grad_norm")
        if loss is not None:
            if isinstance(loss, str):
                try:
                    loss = float(loss)
                except ValueError:
                    pass
            if isinstance(loss, float):
                last_loss = loss
                if math.isnan(loss) or math.isinf(loss) or loss > 100.0:
                    bad_reason = f"loss={loss}"
                    break
        if gn is not None:
            if isinstance(gn, str):
                try:
                    gn = float(gn)
                except ValueError:
                    pass
            if isinstance(gn, float):
                last_gn = gn
                if math.isnan(gn) or math.isinf(gn) or gn > 50.0:
                    bad_reason = f"grad_norm={gn}"
                    break

    if bad_reason is not None:
        print(
            f"[train] !! 학습 발산 감지 ({bad_reason}, last loss={last_loss}, "
            f"last grad_norm={last_gn}) — LoRA 저장 SKIP. "
            "기존 어댑터(있다면) 가 그대로 남고, 추론은 영향을 받지 않습니다.",
            flush=True,
        )
        # 명시적인 실패 코드 — 자동학습 워커가 캐시를 무효화하지 않도록.
        sys.exit(2)

    model.save_pretrained(args.output_dir)
    processor.save_pretrained(args.output_dir)
    print(
        f"[train] LoRA adapter saved → {args.output_dir} "
        f"(last loss={last_loss}, last grad_norm={last_gn})"
    )


if __name__ == "__main__":
    main()