Gemma 4 の複数 LoRA を組み合わせて使う — マージ技法と動的スイッチングで『1モデルで何でもこなす』を本番で実現する

# requirements.txt の中身（バージョンは自分の GPU とドライバで必ず確かめてください）
# torch
# transformers   # Gemma 4 に対応した最新版
# peft
# accelerate
# safetensors
# httpx fastapi uvicorn   # 後半のサーバーとベンチマーク用
 
python -m venv .venv && source .venv/bin/activate
pip install -U -r requirements.txt
 
# Hugging Face にログイン（Gemma 4 の重みを取得するため）
hf auth login
 
# 動作確認: Gemma 4 E2B の最小ロードと1回の生成
python - <<'PY'
from transformers import AutoProcessor, AutoModelForMultimodalLM
 
MODEL_ID = "google/gemma-4-E2B-it"
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = AutoModelForMultimodalLM.from_pretrained(MODEL_ID, dtype="auto", device_map="auto")
 
messages = [{"role": "user", "content": "次の文を一行で要約してください: 壁紙アプリの更新後、読み込みが遅いというレビューが三件続きました。"}]
inputs = processor.apply_chat_template(
    messages, tokenize=True, return_dict=True, return_tensors="pt",
    add_generation_prompt=True, enable_thinking=False,
).to(model.device)
n = inputs["input_ids"].shape[-1]
out = model.generate(**inputs, max_new_tokens=64, do_sample=False)
print(processor.decode(out[0][n:], skip_special_tokens=True))
PY

# lora_preflight.py
# 役割: 複数 LoRA の adapter_config.json を読み、マージ可能かを事前に判定する。
# 使い方: python lora_preflight.py ./artifacts/lora_summary ./artifacts/lora_translate ...
# 終了コード: 0=そのままマージ可 / 1=手法の選び直しが必要 / 2=マージ不可（土台違い）
import json
import sys
from pathlib import Path
 
def load(path: str) -> dict:
    cfg = json.loads((Path(path) / "adapter_config.json").read_text())
    return {
        "path": path,
        "base": cfg.get("base_model_name_or_path"),
        "r": cfg.get("r"),
        "alpha": cfg.get("lora_alpha"),
        "targets": sorted(cfg.get("target_modules") or []),
    }
 
def main(paths: list[str]) -> int:
    cfgs = [load(p) for p in paths]
    for c in cfgs:
        print(f"{c['path']}: base={c['base']} r={c['r']} alpha={c['alpha']} targets={len(c['targets'])}")
 
    if len({c["base"] for c in cfgs}) > 1:
        print("❌ 土台モデルが一致しません。同じ土台で学習し直すか、マージ対象から外してください")
        return 2
 
    code = 0
    if len({c["r"] for c in cfgs}) > 1:
        print("⚠️ ランクが揃っていません。linear / ties は使えないため、cat か *_svd 系を選んでください")
        code = 1
    if len({tuple(c["targets"]) for c in cfgs}) > 1:
        print("⚠️ 対象層が異なります。片方にしか無い層は、その LoRA の効果だけが残ります")
        code = max(code, 1)
    if len({c["alpha"] / c["r"] for c in cfgs if c["r"]}) > 1:
        print("⚠️ alpha/r の比が異なります。同じ重みでも効き方が変わるため、weight 側で補正してください")
        code = max(code, 1)
    if code == 0:
        print("✅ 土台・ランク・対象層が一致しています。linear / ties をそのまま試せます")
    return code
 
if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))

# merge_weighted.py
# 役割: 要約用 LoRA と翻訳用 LoRA を重み付きでマージし、単一モデルとして保存する。
# 前提: lora_preflight.py が 0 を返していること（ランクと土台が一致）。
 
from pathlib import Path
from peft import PeftModel
from transformers import AutoModelForMultimodalLM, AutoProcessor
 
BASE = "google/gemma-4-E2B-it"
LORA_A = "./artifacts/lora_summary"     # 事前に学習済み
LORA_B = "./artifacts/lora_translate"   # 事前に学習済み
OUT_DIR = Path("./artifacts/merged_weighted_6_4")
WEIGHTS = {"summary": 0.6, "translate": 0.4}
 
def merge_weighted():
    processor = AutoProcessor.from_pretrained(BASE)
    base = AutoModelForMultimodalLM.from_pretrained(BASE, dtype="auto", device_map="auto")
 
    # 複数アダプタを名前付きでロード
    model = PeftModel.from_pretrained(base, LORA_A, adapter_name="summary")
    model.load_adapter(LORA_B, adapter_name="translate")
 
    # 重み付き線形結合。ランクが違う場合は combination_type="cat" に変える
    model.add_weighted_adapter(
        adapters=list(WEIGHTS.keys()),
        weights=list(WEIGHTS.values()),
        adapter_name="merged",
        combination_type="linear",
    )
    model.set_adapter("merged")  # これを忘れると別のアダプタが焼き込まれる
 
    merged = model.merge_and_unload()
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    merged.save_pretrained(OUT_DIR, safe_serialization=True)
    processor.save_pretrained(OUT_DIR)
    print(f"✅ saved to {OUT_DIR}")
 
if __name__ == "__main__":
    merge_weighted()

# merge_ties_peft.py（merge_weighted.py からの差分部分）
ADAPTERS = {
    "summary":     ("./artifacts/lora_summary",     0.3),
    "translate":   ("./artifacts/lora_translate",   0.3),
    "code_review": ("./artifacts/lora_code_review", 0.2),
    "sql":         ("./artifacts/lora_sql",         0.2),
}
 
first, (path0, _) = next(iter(ADAPTERS.items()))
model = PeftModel.from_pretrained(base, path0, adapter_name=first)
for name, (path, _) in list(ADAPTERS.items())[1:]:
    model.load_adapter(path, adapter_name=name)
 
model.add_weighted_adapter(
    adapters=list(ADAPTERS.keys()),
    weights=[w for _, w in ADAPTERS.values()],
    adapter_name="merged_ties",
    combination_type="ties",      # ランク違いなら "ties_svd"
    density=0.7,                  # 上位 70% の変化量だけを残す
    majority_sign_method="total",
)
model.set_adapter("merged_ties")
merged = model.merge_and_unload()

# merge_ties.yaml
# 実行: mergekit-yaml merge_ties.yaml ./artifacts/merged_ties_4 --cuda
# 注意: お使いの mergekit が Gemma 4 のアーキテクチャに対応しているかを先に確かめてください
models:
  - model: google/gemma-4-E2B-it+./artifacts/lora_summary
    parameters: { density: 0.7, weight: 0.3 }
  - model: google/gemma-4-E2B-it+./artifacts/lora_translate
    parameters: { density: 0.7, weight: 0.3 }
  - model: google/gemma-4-E2B-it+./artifacts/lora_code_review
    parameters: { density: 0.6, weight: 0.2 }
  - model: google/gemma-4-E2B-it+./artifacts/lora_sql
    parameters: { density: 0.6, weight: 0.2 }
merge_method: ties
base_model: google/gemma-4-E2B-it
parameters:
  normalize: true
  int8_mask: true
dtype: bfloat16

# serve_dynamic.py
# 役割: Gemma 4 本体は1つだけロードし、リクエストの task に応じて LoRA を切り替えて推論する。
# 注意: 切り替えは排他的。並列リクエストは asyncio.Lock で直列化する。
 
import asyncio
from contextlib import asynccontextmanager
 
import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from peft import PeftModel
from transformers import AutoModelForMultimodalLM, AutoProcessor
 
BASE = "google/gemma-4-E2B-it"
ADAPTERS = {
    "summary":     "./artifacts/lora_summary",
    "translate":   "./artifacts/lora_translate",
    "code_review": "./artifacts/lora_code_review",
    "sql":         "./artifacts/lora_sql",
}
 
state = {"model": None, "proc": None, "current": None,
         "lock": asyncio.Lock(), "switches": 0, "requests": 0}
 
@asynccontextmanager
async def lifespan(app: FastAPI):
    proc = AutoProcessor.from_pretrained(BASE)
    base = AutoModelForMultimodalLM.from_pretrained(BASE, dtype="auto", device_map="auto")
    first_name, first_path = next(iter(ADAPTERS.items()))
    model = PeftModel.from_pretrained(base, first_path, adapter_name=first_name)
    for name, path in list(ADAPTERS.items())[1:]:
        model.load_adapter(path, adapter_name=name)
    model.set_adapter(first_name)
    model.eval()
    state.update(model=model, proc=proc, current=first_name)
    print(f"✅ loaded {len(ADAPTERS)} adapters, active={first_name}")
    yield
 
app = FastAPI(lifespan=lifespan)
 
class Req(BaseModel):
    task: str
    prompt: str
    max_new_tokens: int = 256
 
@app.post("/generate")
async def generate(req: Req):
    if req.task not in ADAPTERS:
        raise HTTPException(400, f"unknown task: {req.task}")
    async with state["lock"]:
        state["requests"] += 1
        if state["current"] != req.task:
            state["model"].set_adapter(req.task)
            state["current"] = req.task
            state["switches"] += 1
        proc = state["proc"]
        inputs = proc.apply_chat_template(
            [{"role": "user", "content": req.prompt}],
            tokenize=True, return_dict=True, return_tensors="pt",
            add_generation_prompt=True, enable_thinking=False,
        ).to(state["model"].device)
        n = inputs["input_ids"].shape[-1]
        try:
            with torch.inference_mode():
                out = state["model"].generate(**inputs, max_new_tokens=req.max_new_tokens, do_sample=False)
        except torch.cuda.OutOfMemoryError as e:
            raise HTTPException(503, "GPU OOM — reduce max_new_tokens") from e
        return {"task": req.task, "text": proc.decode(out[0][n:], skip_special_tokens=True)}
 
@app.get("/stats")
async def stats():
    r = state["requests"] or 1
    return {"requests": state["requests"], "switches": state["switches"],
            "switch_ratio": round(state["switches"] / r, 3)}
 
# 起動: uvicorn serve_dynamic:app --port 8000

# bench_switch_churn.py
# 役割: serve_dynamic.py に対し、タスクの並び順だけを変えて負荷をかけ、P50 / P99 / 切り替え率を比べる。
# 前提: golden_prompts.json に {"summary": [...], "translate": [...], ...} の形で評価用プロンプトを置く。
import asyncio
import json
import random
import time
 
import httpx
 
URL = "http://localhost:8000"
TASKS = ["summary", "translate", "code_review", "sql"]
PROMPTS = json.load(open("golden_prompts.json", encoding="utf-8"))
 
def make_order(n: int, mode: str) -> list[str]:
    if mode == "interleaved":            # 毎回タスクが変わる最悪ケース
        return [TASKS[i % len(TASKS)] for i in range(n)]
    block = n // len(TASKS)              # タスクごとにまとめて届く最良ケース
    return [t for t in TASKS for _ in range(block)]
 
def pct(values: list[float], q: float) -> float:
    s = sorted(values)
    return s[min(len(s) - 1, int(q * len(s)))]
 
async def run(mode: str, n: int = 200, concurrency: int = 8) -> None:
    order = make_order(n, mode)
    sem = asyncio.Semaphore(concurrency)
    lat: list[float] = []
    async with httpx.AsyncClient(timeout=180) as c:
        before = (await c.get(f"{URL}/stats")).json()
 
        async def one(task: str) -> None:
            async with sem:
                t0 = time.perf_counter()
                r = await c.post(f"{URL}/generate", json={
                    "task": task, "prompt": random.choice(PROMPTS[task]), "max_new_tokens": 128})
                r.raise_for_status()
                lat.append((time.perf_counter() - t0) * 1000)
 
        t0 = time.perf_counter()
        await asyncio.gather(*(one(t) for t in order))
        wall = time.perf_counter() - t0
        after = (await c.get(f"{URL}/stats")).json()
 
    switches = after["switches"] - before["switches"]
    print(f"{mode:12s} n={n} rps={n / wall:5.2f} "
          f"p50={pct(lat, .50):6.0f}ms p99={pct(lat, .99):6.0f}ms switches={switches}")
 
async def main() -> None:
    random.seed(7)
    for mode in ("grouped", "interleaved"):
        await run(mode)
 
if __name__ == "__main__":
    asyncio.run(main())