예제 소스
오목 좌표 LoRA: 15×15와 19×19 격자를 정확히 읽히기 · 파일 59개
train_lora_mlx.py
"""mlx-vlm 의 LoRA 파인튜닝 CLI 래퍼.
배경 모델을 우리 다운로드 캐시(`vlm_models/`) 에서 찾고, 우리 데이터셋
(`sft_data/baduk_sft/train.jsonl`) 을 입력으로 받아 학습합니다.
mlx-vlm 0.4.x 의 `lora.py` CLI 가 `load_dataset(args.dataset, ...)` 식으로
dataset 을 읽기 때문에 jsonl 을 직접 받지 못합니다. 이 래퍼는 다음을 자동 수행:
1) train.jsonl → HF parquet 변환 (`<output-dir>/_hf_dataset/`)
2) `python -m mlx_vlm.lora --dataset <_hf_dataset>` 으로 학습 실행
사용 예:
python -m vlm_server.train_lora \
--dataset-jsonl sft_data/baduk_sft/train.jsonl \
--model mlx-community/gemma-4-e4b-it-4bit \
--output-dir vlm_models/lora_baduk \
--epochs 1 --batch-size 1 --grad-accum 8 --lr 2e-4 \
--max-samples 200
"""
from __future__ import annotations
import argparse
import json
import os
import shlex
import shutil
import subprocess
import sys
from pathlib import Path
from typing import List, Optional
def _adapter_dir_for_mlx_vlm(path: Optional[str]) -> Optional[str]:
"""mlx_vlm.lora --adapter-path 는 adapter_config.json 이 있는 폴더를 받는다.
사용자는 adapters.safetensors 파일 경로를 넘기기도 하므로, 파일이면 부모 디렉터리로 정규화.
"""
if not path:
return None
p = Path(path).resolve()
if p.is_file() and p.suffix == ".safetensors":
p = p.parent
if not (p / "adapter_config.json").is_file():
raise SystemExit(
f"LoRA 재개 경로에 adapter_config.json 이 없습니다: {p}\n"
f" 예: --resume-from vlm_models/lora_baduk_v2\n"
f" (또는 .../lora_baduk_v2/adapters.safetensors)"
)
return str(p)
def _build_hf_parquet(jsonl_path: Path, dataset_dir: Path, max_samples: int = 0) -> int:
"""jsonl → 단일 parquet 파일로 변환. 반환: 샘플 수."""
try:
from datasets import Dataset
except ImportError as e:
raise SystemExit("`datasets` 패키지가 필요합니다. 의존성을 다시 설치하세요.") from e
print(f"[train_lora] reading {jsonl_path}")
rows: List[dict] = []
with open(jsonl_path, "r", encoding="utf-8") as f:
for line in f:
line = line.strip()
if not line:
continue
try:
rec = json.loads(line)
except json.JSONDecodeError:
continue
# mlx-vlm 가 기대하는 컬럼: messages (필수), images (선택, gemma 는 messages 안에 image 도 들어감)
row = {
"messages": rec.get("messages"),
"images": rec.get("images") or [],
}
if not row["messages"]:
continue
rows.append(row)
if max_samples and len(rows) >= max_samples:
break
if not rows:
raise SystemExit(f"jsonl 에서 유효한 레코드가 없습니다: {jsonl_path}")
print(f"[train_lora] {len(rows)} samples → parquet")
dataset_dir.mkdir(parents=True, exist_ok=True)
# 기존 *.parquet 정리해 두면 load_dataset 가 깔끔히 인식
for old in dataset_dir.glob("*.parquet"):
old.unlink()
ds = Dataset.from_list(rows)
out_pq = dataset_dir / "train.parquet"
ds.to_parquet(str(out_pq))
print(f"[train_lora] wrote {out_pq} ({len(rows)} rows)")
return len(rows)
def main(argv: Optional[List[str]] = None) -> int:
p = argparse.ArgumentParser(description="Gemma 4 LoRA 학습 래퍼 (mlx-vlm 0.4.x)")
p.add_argument("--dataset-jsonl", required=True, help="dataset_baduk.py 가 만든 train.jsonl")
p.add_argument(
"--model",
default="********",
help="HF 모델 ID 또는 로컬 경로 (mlx-vlm 호환).",
)
p.add_argument(
"--output-dir",
default="vlm_models/lora_baduk",
help="LoRA 어댑터 + 파생 데이터셋 출력 폴더",
)
p.add_argument("--epochs", type=int, default=1)
p.add_argument("--batch-size", type=int, default=1)
p.add_argument(
"--grad-accum",
type=int,
default=8,
help="mlx_vlm --gradient-accumulation-steps 값",
)
p.add_argument("--lr", type=float, default=2e-4)
p.add_argument("--lora-rank", type=int, default=8)
p.add_argument("--lora-alpha", type=float, default=16.0)
p.add_argument("--lora-dropout", type=float, default=0.0)
p.add_argument(
"--max-samples",
type=int,
default=0,
help="0=전체. MPS 메모리 부족하면 200~500 정도로 시작.",
)
p.add_argument(
"--max-seq-length",
type=int,
default=2048,
help="모델에 들어갈 최대 토큰 길이.",
)
p.add_argument(
"--grad-checkpoint",
action="store_true",
help="gradient checkpointing 켜기 (메모리 절약, 학습 약간 느려짐).",
)
p.add_argument(
"--resume-from",
default=None,
help="이전 LoRA 폴더 또는 adapters.safetensors 경로 (mlx_vlm 은 폴더를 요구).",
)
p.add_argument(
"--val-jsonl",
default=None,
help="검증용 jsonl. 있으면 별도 parquet 으로 변환되어 mlx_vlm 에 전달.",
)
p.add_argument(
"--mlx-vlm-extra",
nargs=argparse.REMAINDER,
default=[],
help="이 뒤로 오는 모든 인자는 mlx_vlm.lora 에 그대로 전달됩니다.",
)
args = p.parse_args(argv)
ds_path = Path(args.dataset_jsonl).resolve()
if not ds_path.exists():
# 친절한 에러: 실제로 쓸 수 있는 train.jsonl 후보를 훑어서 보여준다.
print(f"데이터셋이 없습니다: {ds_path}", file=sys.stderr)
workspace_root = Path(__file__).resolve().parent
sft_root = workspace_root / "sft_data"
candidates: List[Path] = []
if sft_root.exists():
for p in sorted(sft_root.glob("*/train.jsonl")):
try:
n_lines = sum(1 for _ in open(p, "r", encoding="utf-8"))
except OSError:
n_lines = -1
candidates.append(p)
print(
f" - 사용 가능: {p} (약 {n_lines} 샘플)",
file=sys.stderr,
)
if not candidates:
print(
" (sft_data/ 에 train.jsonl 이 하나도 없습니다. 먼저 "
"'python -m vlm_server.dataset_baduk ...' 으로 데이터셋을 만드세요.)",
file=sys.stderr,
)
else:
pick = candidates[-1]
print(
"\n예시:\n python -m vlm_server.train_lora \\\n"
f" --dataset-jsonl {pick.relative_to(workspace_root)} \\\n"
f" --val-jsonl {pick.parent.relative_to(workspace_root)}/val.jsonl \\\n"
" --model mlx-community/gemma-4-e4b-it-4bit \\\n"
f" --output-dir vlm_models/lora_{pick.parent.name} \\\n"
" --epochs 3 --batch-size 1 --grad-accum 16 --lr 1e-4 \\\n"
" --lora-rank 16 --lora-alpha 32 --grad-checkpoint",
file=sys.stderr,
)
return 2
out_dir = Path(args.output_dir).resolve()
out_dir.mkdir(parents=True, exist_ok=True)
# 1) jsonl → parquet (train + 옵션 val)
hf_dataset_dir = out_dir / "_hf_dataset"
n = _build_hf_parquet(ds_path, hf_dataset_dir, max_samples=args.max_samples)
if n == 0:
return 3
if args.val_jsonl:
val_ds_path = Path(args.val_jsonl).resolve()
if val_ds_path.exists():
val_dir = out_dir / "_hf_dataset_val"
n_val = _build_hf_parquet(val_ds_path, val_dir, max_samples=0)
print(f"[train_lora] val parquet ready: {n_val} rows")
else:
print(f"[train_lora] WARN: --val-jsonl 가 없음: {val_ds_path}")
# 2) mlx_vlm.lora 호출
#
# 주의: mlx_vlm 0.4.x 의 lora.py 에는 --epochs 계산 버그가 있음:
# iters = (len(dataset) // batch_size) * epochs
# dataset = dataset.select(range(iters)) # ← 범위 초과 IndexError
# 해결: --epochs 를 넘기지 않고 --iters = len(dataset)//batch_size 로 단 epoch 씩
# 호출하고, 우리 래퍼가 epochs 만큼 **subprocess 를 반복** (매번 이전 어댑터 resume).
out_path = out_dir / "adapters.safetensors"
iters_per_epoch = max(1, n // max(1, args.batch_size))
workspace_root = Path(__file__).resolve().parent
cache_dir = workspace_root / "vlm_models"
env = dict(os.environ)
env.setdefault("HF_HOME", str(cache_dir))
env.setdefault("HF_HUB_CACHE", str(cache_dir))
env.setdefault("HF_DATASETS_CACHE", str(cache_dir / "datasets_cache"))
def _one_epoch(epoch_idx: int, resume_from: Optional[str]) -> int:
cmd = [
sys.executable,
"-m",
"mlx_vlm.lora",
"--model-path",
args.model,
"--dataset",
str(hf_dataset_dir),
"--iters",
str(iters_per_epoch),
"--batch-size",
str(args.batch_size),
"--gradient-accumulation-steps",
str(args.grad_accum),
"--learning-rate",
str(args.lr),
"--lora-rank",
str(args.lora_rank),
"--lora-alpha",
str(args.lora_alpha),
"--lora-dropout",
str(args.lora_dropout),
"--max-seq-length",
str(args.max_seq_length),
"--output-path",
str(out_path),
]
if args.grad_checkpoint:
cmd.append("--grad-checkpoint")
# mlx-vlm: --adapter-path = adapter_config.json + adapters.safetensors 가 있는 디렉터리.
adapter_dir = _adapter_dir_for_mlx_vlm(resume_from)
if adapter_dir:
cmd += ["--adapter-path", adapter_dir]
if args.mlx_vlm_extra:
cmd += list(args.mlx_vlm_extra)
print(
f"[train_lora] epoch {epoch_idx + 1}/{args.epochs} "
f"iters={iters_per_epoch} dataset_size={n}"
)
print("[train_lora] running:", " ".join(shlex.quote(x) for x in cmd))
return subprocess.call(cmd, env=env)
# 첫 epoch: --resume-from 인자가 있으면 그걸로. 이후 epoch: output-dir(어댑터 폴더) resume.
resume_target: Optional[str] = args.resume_from
rc = 0
for epoch in range(args.epochs):
rc = _one_epoch(epoch, resume_target)
if rc != 0:
print(
f"[train_lora] mlx_vlm.lora 가 epoch {epoch + 1} 에서 실패 (rc={rc})",
file=sys.stderr,
)
return rc
# 다음 epoch 부터는 방금 저장한 어댑터 폴더를 이어받아 학습.
resume_target = str(out_dir)
print(
f"[train_lora] DONE. 어댑터: {out_path}\n"
f" 서버 적용: --lora-dir {out_dir}"
)
return rc
if __name__ == "__main__":
sys.exit(main())