From 0af616231623294198b74fb1bbaf3c71fae181d0 Mon Sep 17 00:00:00 2001 From: dualchimerra Date: Sun, 20 Sep 2026 14:24:17 +0200 Subject: [PATCH] feat(training): enable LyCORIS 4 fused kernels --- docs/adr/0014-lycoris4-fused-kernels.md | 64 +++++++++++++ docs/adr/README.md | 1 + docs/user-guide/training-tips.md | 20 ++++ requirements.txt | 9 +- studio/cli.py | 120 ++++++++++++++++++++++++ tests/test_lycoris_patch.py | 37 ++++---- tests/test_lycoris_tlora.py | 116 ++++++++++++++--------- tests/test_studio_cli.py | 110 ++++++++++++++++++++++ tools/launcher.py | 10 +- utils/lycoris_adapter.py | 102 +++++++++++++++++--- utils/lycoris_patch.py | 54 ++++++----- 11 files changed, 544 insertions(+), 99 deletions(-) create mode 100644 docs/adr/0014-lycoris4-fused-kernels.md diff --git a/docs/adr/0014-lycoris4-fused-kernels.md b/docs/adr/0014-lycoris4-fused-kernels.md new file mode 100644 index 0000000..8c435bd --- /dev/null +++ b/docs/adr/0014-lycoris4-fused-kernels.md @@ -0,0 +1,64 @@ +# 0014 — LyCORIS 4 fused kernels 与 Windows Triton + +**状态**：Accepted + +**日期**：2026-09-20 +**决策者**：项目维护者 + +## 背景 + +训练栈通过 `utils/lycoris_adapter.py` 接入 LyCORIS 3.4。普通 LoRA 已走 +bypass forward，但 LoHa / LoKr / DoRA 与 plain T-LoRA 的 rebuild 路径仍可能 +物化完整 `ΔW`。LyCORIS 4.0 新增 Triton / TileLang fused kernels，并按 +Triton → TileLang → `torch.compile` → eager 的顺序逐调用选择与回退。 + +Windows 上 PyTorch 不自带可导入的 Triton，缺失时还会让 xformers 在启动期打印 +非致命 traceback。当前项目的目标本机环境是 PyTorch 2.11 + CUDA 12.8；该组合 +按 triton-windows 的兼容表对应 Triton 3.6。 + +同时，LyCORIS 4.0 移除了旧 patch 所依赖的内部 `make_kron` / +`rebuild_tucker` import，但 4.0.0 的 LoKr `rank_dropout` 仍会在 CPU 创建 mask。 + +## 候选方案 + +1. **保留 LyCORIS 3.4**：风险最低，但无法使用 fused kernels，启动警告保留。 +2. **直接升级并删除本地 adapter / patch**：代码最少，但丢失 family preset、 + checkpoint metadata、训练 hook、T-LoRA timestep mask，并重新暴露 LoKr device bug。 +3. **升级到 4.0，保留薄 integration adapter，数学运算委托官方 functional API**： + 需要小范围兼容改动，但保留项目协议并获得官方 kernel dispatch。 + +## 决策 + +- 固定 `lycoris-lora==4.0.0`，避免未来 major API 静默漂移。 +- Windows dependency 声明允许 Triton 3.6–3.8；启动器在导入 xformers/flash-attn + 之前按已安装 PyTorch minor 自动收窄或修复版本（2.10/2.11→3.6、 + 2.12/2.13→3.7、2.14→3.8）。 +- 保留 `LycorisAdapter`：它只负责 AdapterProtocol、模型族 target preset、保存元数据、 + sample/eval 状态和 T-LoRA step hook；LoRA 数学继续交给 LyCORIS。 +- plain T-LoRA 在两个低秩因子上应用 timestep mask。forward 调用官方 + `functional.locon.bypass_forward_diff`，不物化 `ΔW`；merge/save 路径调用 + `functional.locon.diff_weight`。两者都进入 LyCORIS 4 fused/compile/eager dispatch。 +- LoKr device patch 改为包装 upstream `get_weight`：只临时关闭坏的 dropout 分支， + 调用 upstream（保留 fused dispatch），然后在 `weight.device` 重放 dropout。 +- OrthoLoRA 保持项目自实现。LyCORIS 4 没有与其 Cayley/SVD 参数化等价的官方 + adapter，不能为了去掉 wrapper 改变算法或 checkpoint 语义。 + +## 理由 + +方案 3 把外部算法实现尽量留给上游，同时保留 Anima/Krea2 集成必需的边界。 +显式 pin 能让 device patch 的适用版本可审计。Triton 缺少或某个 shape 不受支持时， +LyCORIS 自带回退，因此加速失败不会改变训练正确性。 + +## 后果 + +- Windows 首次安装多下载约 50 MB，并在首次遇到新 shape 时支付 JIT/tuning 开销。 +- fused kernel 的 microbenchmark 提升不等于整步训练等比例提升；DiT 主干仍占大头。 +- `rank_dropout`、卷积、过大 rank 等超出 fused scope 的调用会自动回退。 +- Triton 与 PyTorch minor 有配对约束；支持表之外的新 PyTorch 会保守回退并要求 + 更新 `_WINDOWS_TRITON_FOR_TORCH`，不会猜测 ABI。 + +## 参考 + +- [LyCORIS 4.0 fused kernels](https://github.com/KohakuBlueleaf/LyCORIS/blob/v4.0.0/docs/kernels/README.md) +- [LyCORIS backend selection](https://github.com/KohakuBlueleaf/LyCORIS/blob/v4.0.0/docs/kernels/backends.md) +- [triton-windows PyTorch compatibility](https://github.com/triton-lang/triton-windows#3-pytorch) diff --git a/docs/adr/README.md b/docs/adr/README.md index 95e5fe9..658adc3 100644 --- a/docs/adr/README.md +++ b/docs/adr/README.md @@ -16,6 +16,7 @@ | 0008 | [studio/ 4 层重构（0.11.0）](0008-studio-restructure-0.11.0.md) | Accepted | 2026-05-28 | | 0009 | [统一日志 + 错误体系（0.12.0）](0009-logging-error-system.md) | Accepted | 2026-05-28 | | 0010 | [preprocess scope 从项目级 download 下沉到 version 级 train](0010-preprocess-train-scope.md) | Accepted | 2026-06-03 | +| 0014 | [LyCORIS 4 fused kernels 与 Windows Triton](0014-lycoris4-fused-kernels.md) | Accepted | 2026-09-20 | ## 状态值 diff --git a/docs/user-guide/training-tips.md b/docs/user-guide/training-tips.md index 8f69085..924da6c 100644 --- a/docs/user-guide/training-tips.md +++ b/docs/user-guide/training-tips.md @@ -493,6 +493,26 @@ masterpiece, best quality, newest, safe, ## 硬件优化 +### LyCORIS 4 fused kernels（Windows / NVIDIA） + +项目使用 LyCORIS 4.0，并在 Windows 自动校正与当前 PyTorch 配对的 Triton； +例如 PyTorch 2.11 对应 `triton-windows 3.6`。启动时应看到类似： + +```text +[studio] LoRA kernels：LyCORIS 4.0.0 · preferred=triton · fused=triton · Triton 3.6.0.post26 +``` + +backend 按调用自动选择：Triton → TileLang → `torch.compile` → eager PyTorch。 +某个 shape、dtype 或 dropout 不适合 fused kernel 时会自动回退，不需要手工切换。 +首次遇到新 shape 会有 JIT / tuning 预热，后续从项目内 `.cache/triton` 复用缓存。 + +普通 LoRA / LoHa / LoKr / DoRA 可使用官方 fused 路径；plain T-LoRA 保留 timestep +rank mask，并调用不创建独立 `ΔW` 的官方 fused bypass。`T-LoRA + Ortho` 使用项目自己的 Cayley/SVD +参数化，不会伪装成 LyCORIS 算法，因此目前不走这套 fused kernels。 + +这些数字只加速 adapter 运算，不代表整个训练步会同倍数提速；模型 attention / MLP、 +数据和采样通常仍是主耗时。 + ### RTX 3090/4090 (24GB) ```yaml diff --git a/requirements.txt b/requirements.txt index c03335a..0a99194 100644 --- a/requirements.txt +++ b/requirements.txt @@ -13,7 +13,14 @@ wandb>=0.16.0 # 总能装上，而 adamw8bit 是可选优化器而非默认。未装时选它 # 会在训练启动期报可操作错误（optimizers/adamw8bit.py # validate），不会跑到一半才炸。装：pip install bitsandbytes -lycoris-lora>=1.9.0 +# LyCORIS 4 brings fused adapter kernels and keeps the wrapper/module API used +# by this project. Pin the major migration so internal compatibility patches +# remain auditable instead of silently running against a future breaking API. +lycoris-lora==4.0.0 +# Windows PyTorch does not bring Triton itself. The Studio bootstrap narrows +# this broad range to the minor paired with the installed PyTorch (for example, +# torch 2.10/2.11 -> Triton 3.6) before importing flash-attn/xformers. +triton-windows>=3.6.0.post26,<3.9; sys_platform == "win32" prodigyopt>=1.0 # Prodigy 优化器（anima_train.py optimizer=prodigy 时必需） # ProdigyPlusScheduleFree (LoganBooker) — Flux/Qwen-Image LoRA 社区主流，解决 Prodigy 的 # mutation ep 问题。pin >=2.0.0 因为 v1.9.2 ↔ v2.0.0 的 state_dict 不兼容。 diff --git a/studio/cli.py b/studio/cli.py index 18e5e1a..59ffefd 100644 --- a/studio/cli.py +++ b/studio/cli.py @@ -16,6 +16,7 @@ import argparse import hashlib import os +import re import shutil import signal import subprocess @@ -392,6 +393,121 @@ def _try_enable_flash_attn() -> None: ) +def _report_lycoris_kernels() -> None: + """打印 LyCORIS 4 的 adapter-kernel 选择，不触发 JIT 编译。 + + ``available_backends`` 只探测模块是否可导入；真正的 shape/device 选择仍由 + LyCORIS 每次调用完成。这里用 ``preferred`` 而不是 ``active``，避免把尚未 + 首次编译的 Triton kernel 描述成已经执行过。 + """ + try: + from importlib.metadata import PackageNotFoundError, version # noqa: PLC0415 + + lycoris_version = version("lycoris-lora") + except PackageNotFoundError: + return + except Exception as exc: # noqa: BLE001 + print(f"[studio] LyCORIS 状态读取失败（{exc}）", file=sys.stderr) + return + + try: + from lycoris.kernels import available_backends, resolve_backend # noqa: PLC0415 + + available = available_backends() + preferred = resolve_backend() + except Exception as exc: # noqa: BLE001 + print( + f"[studio] 警告：LyCORIS {lycoris_version} kernel backend 初始化失败（{exc}）", + file=sys.stderr, + ) + return + + triton_version = None + for dist_name in ("triton-windows", "triton"): + try: + triton_version = version(dist_name) + break + except PackageNotFoundError: + continue + + fused = ",".join(name for name in available if name in {"triton", "tilelang"}) or "none" + triton_label = f" · Triton {triton_version}" if triton_version else "" + _say( + f"LoRA kernels：LyCORIS {lycoris_version} · preferred={preferred} " + f"· fused={fused}{triton_label}" + ) + + +_WINDOWS_TRITON_FOR_TORCH: dict[tuple[int, int], tuple[int, int]] = { + (2, 10): (3, 6), + (2, 11): (3, 6), + (2, 12): (3, 7), + (2, 13): (3, 7), + (2, 14): (3, 8), +} + + +def _major_minor(raw: str) -> Optional[tuple[int, int]]: + match = re.match(r"^(\d+)\.(\d+)", raw) + return (int(match.group(1)), int(match.group(2))) if match else None + + +def _ensure_windows_triton() -> None: + """Match triton-windows to PyTorch before xformers imports Triton. + + PyTorch and Triton minor releases have an ABI/toolchain pairing. A static + requirement cannot express that dependency, so the launcher repairs a + missing or mismatched Windows wheel from the official compatibility table. + Unsupported/older Torch builds keep LyCORIS' compile/eager fallback. + """ + if sys.platform != "win32": + return + + try: + import torch # noqa: PLC0415 + + torch_minor = _major_minor(torch.__version__) + except Exception: # noqa: BLE001 + return + + target = _WINDOWS_TRITON_FOR_TORCH.get(torch_minor) + if target is None: + if torch_minor is not None: + print( + f"[studio] 警告：暂无 torch {torch_minor[0]}.{torch_minor[1]} " + "对应的 Windows Triton pin；LyCORIS 将自动回退", + file=sys.stderr, + ) + return + + from importlib.metadata import PackageNotFoundError, version # noqa: PLC0415 + + try: + installed_minor = _major_minor(version("triton-windows")) + except PackageNotFoundError: + installed_minor = None + if installed_minor == target: + return + + lower = f"{target[0]}.{target[1]}" + upper = f"{target[0]}.{target[1] + 1}" + requirement = f"triton-windows>={lower},<{upper}" + installed = ( + f"{installed_minor[0]}.{installed_minor[1]}" + if installed_minor is not None + else "not installed" + ) + _say( + f"torch {torch_minor[0]}.{torch_minor[1]} requires Triton {lower} " + f"(current: {installed}); installing compatible wheel..." + ) + if _pip_install([requirement]) != 0: + print( + "[studio] 警告：Triton 安装失败；LyCORIS 将使用 compile/eager fallback", + file=sys.stderr, + ) + + def _check_torch_cuda() -> None: """启动期检查 torch 是否能用 CUDA；CPU-only torch 跑训练 / 出图会极慢。 @@ -637,7 +753,9 @@ def cmd_run(args: argparse.Namespace) -> int: if not getattr(args, 'skip_pending', False): _apply_pending_install() _check_torch_cuda() + _ensure_windows_triton() _try_enable_flash_attn() + _report_lycoris_kernels() _check_onnxruntime() url = f"http://{args.host}:{args.port}/" _say(f"启动后端 → {url}") @@ -692,7 +810,9 @@ def cmd_dev(args: argparse.Namespace) -> int: if not getattr(args, 'skip_pending', False): _apply_pending_install() _check_torch_cuda() + _ensure_windows_triton() _try_enable_flash_attn() + _report_lycoris_kernels() _check_onnxruntime() pg = ProcGroup() diff --git a/tests/test_lycoris_patch.py b/tests/test_lycoris_patch.py index dc7e176..acfdd6b 100644 --- a/tests/test_lycoris_patch.py +++ b/tests/test_lycoris_patch.py @@ -1,7 +1,7 @@ -"""utils.lycoris_patch — lycoris-lora 3.4.0 LokrModule.get_weight device patch。 +"""utils.lycoris_patch — LyCORIS 3.4/4.0 LoKr rank-dropout device patch。 覆盖： -- 命中受影响版本（3.4.0）→ patch 真实生效（torch.rand 调用得到 device=weight.device 的 mask） +- 命中受影响版本 → patch 生效且保留 upstream weight builder - 未装 lycoris → skipped_not_installed - 未知版本 → skipped_version_unknown + warn - 同进程内幂等 → skipped_already_patched @@ -52,15 +52,24 @@ def fresh_patch_module(monkeypatch: pytest.MonkeyPatch): delattr(LokrModule, mod._PATCHED_FLAG) +@pytest.mark.parametrize("affected_version", ["3.4.0", "4.0.0"]) def test_apply_on_known_affected_version_patches_get_weight( - fresh_patch_module, monkeypatch: pytest.MonkeyPatch + fresh_patch_module, monkeypatch: pytest.MonkeyPatch, affected_version: str, ) -> None: - """3.4.0 安装时 patch 应替换 LokrModule.get_weight 且生成的 mask 在 weight.device 上。""" + """受影响版本的 mask 应跟随 weight device，且 upstream builder 仍被调用。""" pytest.importorskip("lycoris.modules.lokr") import torch from lycoris.modules.lokr import LokrModule - monkeypatch.setattr(fresh_patch_module, "version", lambda _: "3.4.0") + upstream_calls: list[object] = [] + + def _fake_upstream(self, shape): + upstream_calls.append(shape) + assert self.rank_dropout == 0.0 + return self.weight.clone() + + monkeypatch.setattr(LokrModule, "get_weight", _fake_upstream) + monkeypatch.setattr(fresh_patch_module, "version", lambda _: affected_version) status = fresh_patch_module.apply_lokr_device_patch() assert status == "applied" @@ -77,22 +86,18 @@ def _spy_rand(*args, **kwargs): monkeypatch.setattr(torch, "rand", _spy_rand) - # 用 mock self（不构造完整 LokrModule，单测只关心 torch.rand 这一行） + # 用 mock self，不依赖不同 LyCORIS 版本的 LoKr constructor 细节。 class _FakeSelf: training = True rank_dropout = 0.5 rank_dropout_scale = False - use_w1 = True - use_w2 = True - tucker = False - scale = 1.0 - lokr_w1 = torch.eye(4) - lokr_w2 = torch.eye(4) + weight = torch.eye(4) fake = _FakeSelf() - # 直接调被替换后的 get_weight；shape=None 让 weight 保持原状 - LokrModule.get_weight(fake, None) - assert captured["device"] is not None, "patched get_weight 没把 device 传给 torch.rand" + LokrModule.get_weight(fake, (4, 4)) + assert upstream_calls == [(4, 4)] + assert fake.rank_dropout == 0.5, "wrapper 必须恢复 module 的 rank_dropout" + assert captured["device"] == fake.weight.device def test_apply_when_not_installed_returns_skipped( @@ -123,6 +128,6 @@ def test_apply_idempotent( ) -> None: """同进程内重复调用：第一次 applied，之后 skipped_already_patched。""" pytest.importorskip("lycoris.modules.lokr") - monkeypatch.setattr(fresh_patch_module, "version", lambda _: "3.4.0") + monkeypatch.setattr(fresh_patch_module, "version", lambda _: "4.0.0") assert fresh_patch_module.apply_lokr_device_patch() == "applied" assert fresh_patch_module.apply_lokr_device_patch() == "skipped_already_patched" diff --git a/tests/test_lycoris_tlora.py b/tests/test_lycoris_tlora.py index 027c263..4bf3a27 100644 --- a/tests/test_lycoris_tlora.py +++ b/tests/test_lycoris_tlora.py @@ -140,6 +140,59 @@ def test_tlora_clear_mask_restores_full_rank() -> None: assert int(adapter._tlora_mask.sum().item()) == min_rank +def test_tlora_weight_rebuild_uses_lycoris4_functional_dispatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Plain T-LoRA keeps its mask but delegates ΔW to LyCORIS 4 kernels.""" + from lycoris.functional import locon as functional_locon + + class _TinyDiT(nn.Module): + def __init__(self) -> None: + super().__init__() + self.q_proj = nn.Linear(16, 16, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.q_proj(x) + + model = _TinyDiT() + adapter = AnimaLycorisAdapter( + preset=ANIMA_PRESET, + algo="tlora", + rank=4, + alpha=4.0, + tlora_min_rank=1, + ) + adapter.inject(model) + adapter._set_tlora_mask(torch.tensor([0.5])) + + calls: list[tuple[torch.Tensor, ...]] = [] + real_diff_weight = functional_locon.diff_weight + + def _spy_diff_weight(*weights, **kwargs): + calls.append(weights) + return real_diff_weight(*weights, **kwargs) + + monkeypatch.setattr(functional_locon, "diff_weight", _spy_diff_weight) + layer = adapter._tlora_modules[0] + weight = layer.make_weight(device=layer.lora_up.weight.device) + + assert calls, "T-LoRA make_weight должен использовать functional dispatch LyCORIS 4" + assert weight.shape == layer.shape + + bypass_calls: list[tuple[torch.Tensor, ...]] = [] + real_bypass = functional_locon.bypass_forward_diff + + def _spy_bypass(*weights, **kwargs): + bypass_calls.append(weights) + return real_bypass(*weights, **kwargs) + + monkeypatch.setattr(functional_locon, "bypass_forward_diff", _spy_bypass) + model(torch.randn(2, 16)) + + assert layer.bypass_mode is True + assert bypass_calls, "T-LoRA forward должен использовать fused bypass API LyCORIS 4" + + # ── test 3: sample_runner 入口调 clear_timestep_mask ──────────────────── @@ -251,46 +304,25 @@ def _noop_ctx(*a, **kw): sr.optimizer_eval_mode = orig_eval # type: ignore[assignment] -# ── test 4: bypass_mode invariant ────────────────────────────────────────── -# -# 设计 invariant：AnimaLycorisAdapter.inject() 只允许两条 bypass 路径： -# 1. self.algo == "lora" and not self.weight_decompose -# 2. self.algo == "lokr" and has_fp8_base -# algo == "tlora" 永远不进这两条分支，必走 rebuild (make_weight) 路径， -# 让 _install_tlora_masks 的 mask patch 真正生效。 -# -# 这里不依赖完整 inject (lycoris preset 与本机版本不兼容)，直接 inspect -# AnimaLycorisAdapter.inject 源码确认 tlora 不会被设 bypass_mode=True。 - - -def test_tlora_inject_never_sets_bypass_mode() -> None: - """inject() 只允许 LoRA 或 FP8 LoKr 设置 bypass_mode=True； - algo='tlora' 不进，必走 make_weight rebuild 路径让 mask patch 生效。 - - 防止后续维护误把 tlora 也加进 bypass 路径让 mask 静默失效。""" - import re - import inspect - src = inspect.getsource(AnimaLycorisAdapter.inject) - # 找形如 `extra["bypass_mode"] = ...` 或 `extra['bypass_mode'] = ...` 的赋值 - pattern = re.compile(r'extra\[["\']bypass_mode["\']\]\s*=', re.MULTILINE) - matches = list(pattern.finditer(src)) - assert matches, ( - "inject 源码里找不到 extra['bypass_mode'] 赋值；" - "源码结构变了，需重写本 invariant 测试或重新评估 bypass_mode 默认行为。" +# ── test 4: DoRA bypass invariant ────────────────────────────────────────── + + +def test_tlora_dora_keeps_rebuild_path() -> None: + """DoRA 依赖完整 weight decomposition，不能切到 bypass forward。""" + + class _TinyDiT(nn.Module): + def __init__(self) -> None: + super().__init__() + self.q_proj = nn.Linear(16, 16, bias=False) + + adapter = AnimaLycorisAdapter( + preset=ANIMA_PRESET, + algo="tlora", + rank=4, + alpha=4.0, + weight_decompose=True, ) - for m in matches: - # 该赋值之前 ~250 字符内必须出现两个受支持守卫之一，且不能提到 tlora。 - # FP8 LoKr 必须同时检查 has_fp8_base，避免普通 LoKr 被误切到 bypass。 - ctx_start = max(0, m.start() - 250) - ctx = src[ctx_start:m.start()] - has_lora_guard = bool(re.search(r'self\.algo\s*==\s*["\']lora["\']', ctx)) - has_fp8_lokr_guard = bool(re.search( - r'self\.algo\s*==\s*["\']lokr["\']\s+and\s+has_fp8_base', ctx, - )) - mentions_tlora_nearby = "tlora" in ctx.lower() - assert (has_lora_guard or has_fp8_lokr_guard) and not mentions_tlora_nearby, ( - f"extra['bypass_mode'] 赋值 (位置 {m.start()}) 上方 250 字符内未见到" - f" LoRA 或 FP8 LoKr 的独立守卫；可能让 tlora / 普通 LoKr 误走" - f" bypass，让 make_weight 路径静默失效。" - f" 上下文：\n{ctx[-200:]}" - ) + adapter.inject(_TinyDiT()) + + assert adapter._tlora_modules + assert all(not layer.bypass_mode for layer in adapter._tlora_modules) diff --git a/tests/test_studio_cli.py b/tests/test_studio_cli.py index 207b9d9..8fd068a 100644 --- a/tests/test_studio_cli.py +++ b/tests/test_studio_cli.py @@ -22,10 +22,19 @@ @pytest.fixture def fake_calls(monkeypatch: pytest.MonkeyPatch) -> list[list[str]]: calls: list[list[str]] = [] + def fake(cmd, **kwargs: Any) -> int: calls.append(list(cmd)) return 0 + monkeypatch.setattr(cli.subprocess, "call", fake) + # Tests using this fixture exercise command dispatch only. Avoid importing + # the real CUDA/Triton stack in those unit tests. + monkeypatch.setattr(cli, "_check_torch_cuda", lambda: None) + monkeypatch.setattr(cli, "_ensure_windows_triton", lambda: None) + monkeypatch.setattr(cli, "_try_enable_flash_attn", lambda: None) + monkeypatch.setattr(cli, "_report_lycoris_kernels", lambda: None) + monkeypatch.setattr(cli, "_check_onnxruntime", lambda: None) return calls @@ -434,6 +443,105 @@ def test_check_torch_cuda_warns_on_cuda_build_but_unavailable( assert "pip install torch" not in out.err +def test_report_lycoris_kernels_prints_preferred_backend( + monkeypatch: pytest.MonkeyPatch, capsys +) -> None: + """启动诊断展示版本、首选 backend 与可用 fused backend。""" + import importlib.metadata + import sys as _sys + import types + + versions = { + "lycoris-lora": "4.0.0", + "triton-windows": "3.6.0.post26", + } + + def _version(name: str) -> str: + if name in versions: + return versions[name] + raise importlib.metadata.PackageNotFoundError(name) + + kernels = types.ModuleType("lycoris.kernels") + kernels.available_backends = lambda: ("triton", "compile", "torch") + kernels.resolve_backend = lambda: "triton" + lycoris = types.ModuleType("lycoris") + lycoris.__path__ = [] # type: ignore[attr-defined] + monkeypatch.setattr(importlib.metadata, "version", _version) + monkeypatch.setitem(_sys.modules, "lycoris", lycoris) + monkeypatch.setitem(_sys.modules, "lycoris.kernels", kernels) + + cli._report_lycoris_kernels() + out = capsys.readouterr() + assert "LyCORIS 4.0.0" in out.out + assert "preferred=triton" in out.out + assert "fused=triton" in out.out + assert "Triton 3.6.0.post26" in out.out + assert out.err == "" + + +def test_report_lycoris_kernels_silent_when_not_installed( + monkeypatch: pytest.MonkeyPatch, capsys +) -> None: + import importlib.metadata + + def _missing(name: str) -> str: + raise importlib.metadata.PackageNotFoundError(name) + + monkeypatch.setattr(importlib.metadata, "version", _missing) + cli._report_lycoris_kernels() + out = capsys.readouterr() + assert out.out == "" + assert out.err == "" + + +def test_ensure_windows_triton_installs_torch_211_pair( + monkeypatch: pytest.MonkeyPatch, capsys +) -> None: + import importlib.metadata + + _install_fake_torch( + monkeypatch, + _FakeTorch(available=True, cuda_build="12.8", version="2.11.0+cu128"), + ) + monkeypatch.setattr(cli.sys, "platform", "win32") + + def _missing(_name: str) -> str: + raise importlib.metadata.PackageNotFoundError + + installs: list[list[str]] = [] + monkeypatch.setattr(importlib.metadata, "version", _missing) + monkeypatch.setattr( + cli, + "_pip_install", + lambda args: installs.append(args) or 0, + ) + + cli._ensure_windows_triton() + + assert installs == [["triton-windows>=3.6,<3.7"]] + assert "torch 2.11 requires Triton 3.6" in capsys.readouterr().out + + +def test_ensure_windows_triton_keeps_matching_version( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import importlib.metadata + + _install_fake_torch( + monkeypatch, + _FakeTorch(available=True, cuda_build="12.8", version="2.11.0+cu128"), + ) + monkeypatch.setattr(cli.sys, "platform", "win32") + monkeypatch.setattr(importlib.metadata, "version", lambda _name: "3.6.0.post26") + monkeypatch.setattr( + cli, + "_pip_install", + lambda _args: pytest.fail("matching Triton must not be reinstalled"), + ) + + cli._ensure_windows_triton() + + # --------------------------------------------------------------------------- # PR-5 — _print_npm_install_hint # --------------------------------------------------------------------------- @@ -541,7 +649,9 @@ def _stub_run_bootstrap( monkeypatch.setattr(cli, "_apply_update_pending", lambda: None, raising=False) monkeypatch.setattr(cli, "_apply_pending_install", lambda: None) monkeypatch.setattr(cli, "_check_torch_cuda", lambda: None) + monkeypatch.setattr(cli, "_ensure_windows_triton", lambda: None) monkeypatch.setattr(cli, "_try_enable_flash_attn", lambda: None) + monkeypatch.setattr(cli, "_report_lycoris_kernels", lambda: None) monkeypatch.setattr(cli, "_check_onnxruntime", lambda: None) monkeypatch.setattr(cli, "_spawn_browser_opener", lambda *a, **k: None) return dist diff --git a/tools/launcher.py b/tools/launcher.py index 6da6c1f..4e975f4 100644 --- a/tools/launcher.py +++ b/tools/launcher.py @@ -531,13 +531,17 @@ def main(argv: Optional[list[str]] = None) -> int: parser = build_parser() args, passthrough = parser.parse_known_args(argv) - say(f"{APP_NAME} launcher") + print() + say("+----------------------------------------------------------+") + say(f"| {APP_NAME:<56} |") + say("+----------------------------------------------------------+") repo = find_repo(args.repo) - say(f"folder: {repo}") + say(f"workspace : {repo}") # 在任何 pip 调用之前 —— 见 apply_local_caches 的说明。 cached = apply_local_caches(repo) if cached: - say(f"caches -> {repo / '.cache'} (set ALS_SYSTEM_CACHES=1 to use system locations)") + say(f"cache : {repo / '.cache'}") + say(" ALS_SYSTEM_CACHES=1 restores system locations") if args.check: return run_check(repo) diff --git a/utils/lycoris_adapter.py b/utils/lycoris_adapter.py index a3c9a3f..2c3fc9d 100644 --- a/utils/lycoris_adapter.py +++ b/utils/lycoris_adapter.py @@ -199,7 +199,7 @@ def inject(self, model: nn.Module) -> dict[str, nn.Module]: # 是 LoRA 论文 + sd-scripts + PEFT 的标准 forward；对外行为完全等价但 ~2× 快。 # DoRA(weight_decompose) 路径数学上必须 rebuild —— lycoris bypass forward 不走 wd # 分支，会让 DoRA 静默失效；这里 guard。参考 lycoris docs/Network-Args.md "Bypass Mode"。 - if self.algo == "lora" and not self.weight_decompose: + if self.algo in {"lora", "tlora"} and not self.weight_decompose: extra["bypass_mode"] = True # Krea2 FP8 checkpoints keep frozen Linear weights in float8 and @@ -306,7 +306,12 @@ def _install_tlora_masks(self) -> None: continue original_make_weight = lora.make_weight - def _make_weight_with_tlora(device=None, *, _lora=lora, _original_make_weight=original_make_weight): + def _make_weight_with_tlora( + device=None, + *, + _lora=lora, + _original_make_weight=original_make_weight, + ): wa = _lora.lora_up.weight.to(device) wb = _lora.lora_down.weight.to(device) mask = getattr(_lora, "_anima_tlora_mask", None) @@ -316,17 +321,22 @@ def _make_weight_with_tlora(device=None, *, _lora=lora, _original_make_weight=or view_down = (-1, 1) + (1,) * max(0, wb.dim() - 2) wa = wa * mask.view(*view_up) wb = wb * mask.view(*view_down) - if getattr(_lora, "tucker", False): - t = _lora.lora_mid.weight - wa_t = wa.view(wa.size(0), -1).transpose(0, 1) - wb_t = wb.view(wb.size(0), -1) - from lycoris.modules.locon import rebuild_tucker - weight = rebuild_tucker(t, wa_t, wb_t) - else: - weight = wa.view(wa.size(0), -1) @ wb.view(wb.size(0), -1) - weight = weight.view(_lora.shape) + # LyCORIS 4 functional API owns eager/compile/Triton/TileLang + # selection. Keeping the timestep mask on the factors and + # delegating the rebuild preserves T-LoRA semantics while + # allowing the new fused merge forward/backward kernels. + from lycoris.functional.locon import diff_weight + + t = ( + _lora.lora_mid.weight.to(device) + if getattr(_lora, "tucker", False) + else None + ) + weight = diff_weight(wb, wa, t, gamma=1.0).view(_lora.shape) if _lora.training and _lora.rank_dropout: - drop = (torch.rand(weight.size(0), device=device) > _lora.rank_dropout).to(weight.dtype) + drop = ( + torch.rand(weight.size(0), device=device) > _lora.rank_dropout + ).to(weight.dtype) drop = drop.view(-1, *[1] * len(weight.shape[1:])) if _lora.rank_dropout_scale: drop /= drop.mean() @@ -335,6 +345,74 @@ def _make_weight_with_tlora(device=None, *, _lora=lora, _original_make_weight=or lora._anima_original_make_weight = original_make_weight lora.make_weight = _make_weight_with_tlora + + if getattr(lora, "bypass_mode", False): + original_bypass_forward_diff = lora.bypass_forward_diff + + def _bypass_forward_diff_with_tlora( + x, + scale=1, + *, + _lora=lora, + ): + wa = _lora.lora_up.weight.to(x) + wb = _lora.lora_down.weight.to(x) + mask = getattr(_lora, "_anima_tlora_mask", None) + if mask is not None: + mask = mask.to(device=x.device, dtype=wa.dtype) + view_up = (1, -1) + (1,) * max(0, wa.dim() - 2) + view_down = (-1, 1) + (1,) * max(0, wb.dim() - 2) + wa = wa * mask.view(*view_up) + wb = wb * mask.view(*view_down) + + if _lora.training and _lora.rank_dropout: + drop = ( + torch.rand(_lora.lora_dim, device=x.device) + > _lora.rank_dropout + ).to(wb.dtype) + if _lora.rank_dropout_scale: + drop /= drop.mean() + wb = wb * drop.view( + -1, + *[1] * (wb.dim() - 1), + ) + + t = ( + _lora.lora_mid.weight.to(x) + if getattr(_lora, "tucker", False) + else None + ) + op = ( + _lora.lora_mid + if getattr(_lora, "tucker", False) + else _lora.lora_down + ) + extra_args = ( + { + "stride": op.stride, + "padding": op.padding, + "dilation": op.dilation, + "groups": op.groups, + } + if _lora.isconv + else {} + ) + from lycoris.functional.locon import bypass_forward_diff + + diff = bypass_forward_diff( + x, + None, + wb, + wa, + t, + gamma=_lora.scale * scale, + extra_args=extra_args, + ) + scalar = _lora.scalar.to(device=diff.device, dtype=diff.dtype) + return _lora.dropout(diff * scalar) + + lora._anima_original_bypass_forward_diff = original_bypass_forward_diff + lora.bypass_forward_diff = _bypass_forward_diff_with_tlora lora._anima_tlora_patched = True logger.info( "T-LoRA timestep rank mask enabled on %s/%s LoRA layers (min_rank=%s, alpha_rank_scale=%s)", diff --git a/utils/lycoris_patch.py b/utils/lycoris_patch.py index 985f998..4ee5a30 100644 --- a/utils/lycoris_patch.py +++ b/utils/lycoris_patch.py @@ -1,4 +1,4 @@ -"""lycoris-lora 3.4.0 的 LokrModule.get_weight rank_dropout device bug patch。 +"""LyCORIS LoKr ``rank_dropout`` device bug compatibility patch. 上游 bug：`torch.rand(weight.size(0))` 没传 `device=`，生成 CPU mask， 与 CUDA weight 相乘时报 device mismatch。仅在 `rank_dropout > 0` 且 @@ -8,12 +8,16 @@ - hijack 只保证 sample/eval 时 network 进 eval 模式（不触发 rank_dropout 分支） - 但用户若配置 `rank_dropout > 0`，正常 training step 仍走 rank_dropout 分支 —— hijack 不覆盖这条路径，仍会撞 bug -- 因此从根上把 LokrModule.get_weight 替换成带 `device=` 的版本 +- 因此包装 ``LokrModule.get_weight``，让 dropout mask 跟随 weight device 版本守卫： - 只对 KNOWN_AFFECTED_VERSIONS 内的版本 patch - 其他版本（包括上游已修的版本）log warn 并跳过；避免覆盖上游已 fix 的实现 -- 上游 fix 后请把对应 KNOWN_AFFECTED_VERSIONS 项删掉 +- 上游 fix 后请把对应 ``KNOWN_AFFECTED_VERSIONS`` 项删掉 + +实现刻意不再导入 LyCORIS 的 ``make_kron`` / ``rebuild_tucker`` 等内部函数： +4.0 把这些实现迁移到了 functional kernel API。包装原方法既保留 4.0 的 fused +kernel dispatch，也让补丁跨 3.4/4.0 的内部重构保持稳定。 上游 issue：https://github.com/KohakuBlueleaf/LyCORIS/issues —— 待提 """ @@ -26,9 +30,9 @@ logger = logging.getLogger(__name__) # 已知确认受 rank_dropout device bug 影响的 lycoris-lora 版本。 -# 经实测：3.4.0 的 `lycoris/modules/lokr.py:get_weight` 走 +# 经实测：3.4.0 / 4.0.0 的 `lycoris/modules/lokr.py:get_weight` 走 # `torch.rand(weight.size(0))`（CPU mask），与 CUDA weight 相乘失败。 -KNOWN_AFFECTED_VERSIONS: frozenset[str] = frozenset({"3.4.0"}) +KNOWN_AFFECTED_VERSIONS: frozenset[str] = frozenset({"3.4.0", "4.0.0"}) PatchStatus = Literal[ "applied", # 命中受影响版本，已 patch @@ -51,7 +55,7 @@ def apply_lokr_device_patch() -> PatchStatus: return "skipped_not_installed" try: - from lycoris.modules.lokr import LokrModule, make_kron, rebuild_tucker + from lycoris.modules.lokr import LokrModule except Exception as exc: # pragma: no cover - 装了 lycoris-lora 但 import 异常的边界 logger.warning( "lycoris-lora %s 已安装但 lycoris.modules.lokr 导入失败: %s；跳过 device patch", @@ -74,27 +78,27 @@ def apply_lokr_device_patch() -> PatchStatus: import torch # noqa: PLC0415 延迟到此处避免顶层 import 副作用 + original_get_weight = LokrModule.get_weight + def _get_weight_fixed(self, shape): # type: ignore[no-untyped-def] - weight = make_kron( - self.lokr_w1 if self.use_w1 else self.lokr_w1_a @ self.lokr_w1_b, - ( - self.lokr_w2 - if self.use_w2 - else ( - rebuild_tucker(self.lokr_t2, self.lokr_w2_a, self.lokr_w2_b) - if self.tucker - else self.lokr_w2_a @ self.lokr_w2_b - ) - ), - self.scale, - ) - dtype = weight.dtype - if shape is not None: - weight = weight.view(shape) - if self.training and self.rank_dropout: + # Ask upstream to build the weight (and, in 4.x, select its fused + # kernel) with only the broken dropout branch temporarily disabled. + # Restoring in ``finally`` keeps module state intact even if upstream + # raises for an unsupported shape. + rank_dropout = self.rank_dropout + apply_dropout = bool(self.training and rank_dropout) + if apply_dropout: + self.rank_dropout = 0.0 + try: + weight = original_get_weight(self, shape) + finally: + if apply_dropout: + self.rank_dropout = rank_dropout + + if apply_dropout: drop = ( - torch.rand(weight.size(0), device=weight.device) > self.rank_dropout - ).to(dtype) + torch.rand(weight.size(0), device=weight.device) > rank_dropout + ).to(weight.dtype) drop = drop.view(-1, *[1] * len(weight.shape[1:])) if self.rank_dropout_scale: drop /= drop.mean()