From 919097c73272fcff5804a951a3c13bad527dcfcc Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 17 Oct 2026 19:57:06 +0000 Subject: [PATCH] LoRA training: our own loop for H3, and a Log that survives a crash The H3 run trains with a loop of our own (LoRA on the attention and feed-forward projections of the 50 blocks, per-block checkpoint, uniform noise levels, musubi-tuner's approach) while ComfyUI carries the batches. Every Log record is flushed to disk as it is said and reloaded on start; the run writes each stage and step to it. ComfyUI's own offloading toggle is ignored by the new loop. UNTESTED on the real weights. Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01MpnF7GwNGMbngU11wjym7P --- .gitignore | 1 + core/config.py | 3 + core/log.py | 48 +++++++ core/tests/test_log.py | 16 +++ modules/models/minimax_h3/training.py | 10 ++ modules/training/lora/jobs.py | 14 ++- modules/training/lora/loop.py | 172 ++++++++++++++++++++++++++ modules/training/lora/train.py | 31 +++-- 8 files changed, 282 insertions(+), 13 deletions(-) create mode 100644 modules/training/lora/loop.py diff --git a/.gitignore b/.gitignore index def5b2d2..b59db12a 100644 --- a/.gitignore +++ b/.gitignore @@ -49,3 +49,4 @@ app/playwright-report/ /block_influence.latest.json /phrase_probe.enabled /phrase_probe.latest.json +/log.jsonl diff --git a/core/config.py b/core/config.py index bc1c798f..5ce91772 100644 --- a/core/config.py +++ b/core/config.py @@ -46,3 +46,6 @@ ADDONS_FILE = ROOT / "addons.json" # LoRA training (modules/training/lora): training sets, tagger models, run history. Per machine, like media. TRAINING_DIR = ROOT / "training" + +# The Log's last lines, kept on disk as they are written so a crash does not take them (core/log.py). Per machine. +LOG_FILE = ROOT / "log.jsonl" diff --git a/core/log.py b/core/log.py index 2ab07d7b..34dd5896 100644 --- a/core/log.py +++ b/core/log.py @@ -22,6 +22,7 @@ The buffer is bounded because a long rental session must not turn a log into a memory leak. """ +import json import sys import time from collections import deque @@ -40,6 +41,49 @@ LEVELS = (INFO, ALERT, WARNING, CRITICAL) HISTORY = 500 _records: deque = deque(maxlen=HISTORY) +_file = None # opened on the first record; None while persistence is off (tests, unwritable folder) +_persist = "pytest" not in sys.modules # tests do not write the machine's Log + + +def _path(): + from . import config + return config.LOG_FILE + + +def _load() -> None: + """The last records of the previous session, so a crash leaves the Log as it was. The file is cut back to that many lines.""" + try: + path = _path() + lines = path.read_text(encoding="utf-8").splitlines()[-HISTORY:] + except (OSError, ValueError): + return + for text in lines: + try: + entry = json.loads(text) + if isinstance(entry, dict) and "message" in entry: + _records.append(entry) + except ValueError: + continue + try: + path.write_text("".join(t + "\n" for t in lines), encoding="utf-8") + except OSError: + pass + + +def _keep(entry: dict) -> None: + """One line appended and flushed per record: a dead process leaves everything it said.""" + global _file, _persist + if not _persist: + return + try: + if _file is None: + _file = open(_path(), "a", encoding="utf-8") + _file.write(json.dumps(entry, default=str) + "\n") + _file.flush() + except (OSError, ValueError): + _persist = False # a Log that cannot be written is still shown, just not kept + + # Keys already said once. Cleared per run rather than per process: v4's dedup # lasted the life of the interpreter, so a session reported the first generation # that went inert and stayed quiet for every one after it. @@ -59,6 +103,7 @@ def record(level: str, source: str, message: str) -> dict: entry = {"at": time.time(), "level": level, "source": source, "message": message, "run": run_mod.current()} _records.append(entry) + _keep(entry) print(line(level, source, message), file=sys.stderr, flush=True) return entry @@ -129,3 +174,6 @@ def _reset() -> None: """Tests only.""" _records.clear() _said.clear() + + +_load() diff --git a/core/tests/test_log.py b/core/tests/test_log.py index e6ef29d6..3fb5bf13 100644 --- a/core/tests/test_log.py +++ b/core/tests/test_log.py @@ -140,3 +140,19 @@ def test_different_keys_are_independent(): log.once("a", log.ALERT, "src", "one") log.once("b", log.ALERT, "src", "two") assert len(log.history()) == 2 + + +def test_the_log_survives_the_process_dying(tmp_path, monkeypatch): + """Each record is on disk as soon as it is said; the next start shows the previous session's last lines.""" + from core import config + monkeypatch.setattr(config, "LOG_FILE", tmp_path / "log.jsonl") + monkeypatch.setattr(log, "_persist", True) + monkeypatch.setattr(log, "_file", None) + log._reset() + log.info("training", "step 3: backward") + log.warning("training", "step 4: forward") + log._file.close() + log._records.clear() + monkeypatch.setattr(log, "_file", None) + log._load() + assert [r["message"] for r in log.history()] == ["step 3: backward", "step 4: forward"] diff --git a/modules/models/minimax_h3/training.py b/modules/models/minimax_h3/training.py index 9b94c14b..048ec4b7 100644 --- a/modules/models/minimax_h3/training.py +++ b/modules/models/minimax_h3/training.py @@ -85,6 +85,14 @@ def silence_latent(audio_vae, seconds: float): return audio_vae.encode(wave.movedim(1, -1)) +LORA_LAYERS = ("attn.qkv_proj", "attn.out_proj", "mlp.fc1", "mlp.fc2") + + +def lora_targets(diffusion_model): + """[(name, module)] of the layers a LoRA goes on, named as in the model file (blocks.N.attn.qkv_proj ...).""" + return [(f"blocks.{i}.{layer}", block.get_submodule(layer)) for i, block in enumerate(diffusion_model.blocks) for layer in LORA_LAYERS] + + def adapt(model): """None for any model but H3. For H3, the facts the training module plans around.""" from . import is_h3 @@ -92,6 +100,8 @@ def adapt(model): return None return {"name": "H3", "step": SPATIAL_STEP, "video": True, "batch_limit": 1, "fps": FPS, "blocks": lambda diffusion_model: list(diffusion_model.blocks), # the 50 DiT blocks, for block swap + "lora_targets": lora_targets, # musubi-tuner's choice: the attention and feed-forward projections of the 50 DiT blocks + "checkpoint": lambda diffusion_model: list(diffusion_model.blocks) + list(diffusion_model.token_refiner.blocks), "checkpoint_depth": 2, # one checkpoint per block and one for the text refiner: depth 1 would be the whole model as a single piece "frames": lambda n: align_frame_count(max(5, int(n))), "fit": lambda available: 17 * ((int(available) - 5) // 17) + 5 if available >= 5 else 0, # the longest 17k+5 run that fits in a clip diff --git a/modules/training/lora/jobs.py b/modules/training/lora/jobs.py index 661e9371..b60d7baf 100644 --- a/modules/training/lora/jobs.py +++ b/modules/training/lora/jobs.py @@ -13,7 +13,7 @@ import threading import time import uuid -from ..._core import config +from ..._core import config, log MAX_KEPT_LOSSES = 4000 # a saved run keeps a thinned curve: a 100000 step run is not a 100000 number file _LOCK = threading.Lock() @@ -31,6 +31,17 @@ def _dir(): return path +def trail(message: str) -> None: + """A stage of the run in the Log, which is kept on disk as it is written: when ComfyUI dies mid-run, the last line says where.""" + rss = "" + try: + with open("/proc/self/statm") as f: + rss = f" [process {int(f.read().split()[1]) * 4096 / 1024 ** 3:.1f} GiB in memory]" + except OSError: + pass + log.info("FunPack LoRA training", message + rss) + + def begin(set_name: str, name: str, total: int, settings: dict, prompt_id: str = "") -> str: """Start a record. A run that was still marked running (ComfyUI restarted under it) is simply replaced.""" global _stop @@ -52,6 +63,7 @@ def step(loss: float) -> None: def say(message: str, state: str | None = None) -> None: + trail(message) with _LOCK: _now["message"] = message if state and _now["state"] not in ("done", "error", "stopped"): diff --git a/modules/training/lora/loop.py b/modules/training/lora/loop.py new file mode 100644 index 00000000..ff2e9c60 --- /dev/null +++ b/modules/training/lora/loop.py @@ -0,0 +1,172 @@ +"""FunPack's own LoRA training loop, for models that tell the module what to train (H3 today). + +ComfyUI stays the backend: its model classes and ops, its sampler and model management that carry a +batch through the network. What ComfyUI's trainer decides is done here instead: which layers get a +LoRA, how the blocks are checkpointed, how a step is drawn and weighed, how the optimizer steps. + +The approach follows musubi-tuner by kohya-ss (Apache-2.0): LoRA on the attention and feed-forward +projections of the transformer blocks only, a checkpoint around each block, uniform noise levels, +rank = alpha, the learning rate and optimizer it recommends. The code is written for this module, not copied. +""" + +from __future__ import annotations + +from ..._core import log +from . import jobs + + +class Lora: + """A low-rank pair per target layer, applied by a forward hook so the layer, its ops and its weight casting stay ComfyUI's own.""" + + def __init__(self, targets, rank, alpha, dtype, device, torch): + self.torch, self.rank, self.alpha = torch, int(rank), float(alpha) + self.names = [name for name, _ in targets] + self.pairs, self.handles = {}, [] + nn = torch.nn + for name, module in targets: + weight = module.weight + down = nn.Parameter(torch.empty(self.rank, weight.shape[1:].numel(), dtype=dtype, device=device)) + up = nn.Parameter(torch.zeros(weight.shape[0], self.rank, dtype=dtype, device=device)) + nn.init.kaiming_uniform_(down, a=5 ** 0.5) + self.pairs[name] = (down, up) + self.handles.append(module.register_forward_hook(self._hook(down, up))) + + def _hook(self, down, up): + scale = self.alpha / self.rank + F = self.torch.nn.functional + + def add(module, args, out): + x = args[0] + delta = F.linear(F.linear(x.to(down.dtype), down), up) * scale + return out + delta.to(out.dtype) + return add + + def parameters(self): + return [p for pair in self.pairs.values() for p in pair] + + def state(self, prefix="diffusion_model."): + """The file's tensors, named the way ComfyUI loads a LoRA (lora_down / lora_up / alpha per layer).""" + out = {} + for name, (down, up) in self.pairs.items(): + out[f"{prefix}{name}.lora_down.weight"] = down.detach().to("cpu").contiguous() + out[f"{prefix}{name}.lora_up.weight"] = up.detach().to("cpu").contiguous() + out[f"{prefix}{name}.alpha"] = self.torch.tensor(self.alpha) + return out + + def remove(self): + for handle in self.handles: + handle.remove() + self.handles = [] + + +def checkpoint_modules(modules, torch): + """Each module runs inside a non-reentrant checkpoint: its activations are dropped after the forward and recomputed in the backward, one module at a time. Returns an undo.""" + from torch.utils.checkpoint import checkpoint + undo = [] + for module in modules: + original = module.forward + + def wrapped(*args, __original=original, **kwargs): + def run(args, kwargs): + return __original(*args, **kwargs) + return checkpoint(run, args, kwargs, use_reentrant=False) + module.forward = wrapped + undo.append((module, original)) + + def restore(): + for module, original in undo: + module.forward = original + return restore + + +def optimizer_for(name, params, lr, torch): + """AdamW is musubi's recommendation; the other names are the module's older choices.""" + return {"Adam": torch.optim.Adam, "AdamW": torch.optim.AdamW, "SGD": torch.optim.SGD, "RMSprop": torch.optim.RMSprop}[name](params, lr=lr) + + +class Sampler: + """Handed to ComfyUI's guider in place of a sampler: instead of denoising, it trains. Built lazily so this file imports without comfy.""" + + @staticmethod + def make(nt, torch, comfy_samplers, custom_sampler, lora, optimizer, criterion, cfg, dtype, stopped): + class OwnSampler(comfy_samplers.Sampler): + def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False): + model_wrap.conds = nt.process_cond_list(model_wrap.conds) + cond = model_wrap.conds["positive"] + size = sigmas.size(0) + accumulate, seed = cfg["grad_accumulation"], cfg["seed"] + sampling = model_wrap.inner_model.model_sampling + for i in range(cfg["steps"] * accumulate): + gen = torch.Generator().manual_seed(seed + i) + pick = torch.randperm(size, generator=gen)[:1].tolist() + latent = torch.stack([latent_image[j] for j in pick]) + noise_i = custom_sampler.Noise_RandomNoise(seed + i * 1000).generate_noise({"samples": latent}).to(latent.device) + # uniform noise level, kept off the very ends where the 1/sigma^2 weight below explodes + level = 0.05 + 0.95 * torch.rand((1,), generator=gen).item() + sigma = torch.tensor([level]).to(latent.device) + jobs.trail(f"step {i + 1}: forward (item {pick[0]}, noise level {level:.2f}, card peak {torch.cuda.max_memory_allocated() / 1024 ** 3:.1f} GiB)" if torch.cuda.is_available() else f"step {i + 1}: forward") + xt = sampling.noise_scaling(sigma, noise_i, latent, False) + model_wrap.conds["positive"] = [cond[j] for j in pick] + batch_extra = nt.make_batch_extra_option_dict(extra_args, pick, full_size=size) + with torch.autocast(xt.device.type, dtype=dtype): + pred = model_wrap(xt.requires_grad_(True), sigma.requires_grad_(True), **batch_extra) + # the network answers with the clean latent; weighing its error by 1/sigma^2 makes it the velocity loss musubi trains on + loss = criterion(pred.float(), latent.float()) / (level ** 2) + jobs.trail(f"step {i + 1}: backward") + (loss / accumulate).backward() + value = float(loss.item()) + del pred, xt, loss + jobs.trail(f"step {i + 1}: loss {value:.4f}") + stopped(value) + if (i + 1) % accumulate == 0: + optimizer.step() + optimizer.zero_grad() + torch.cuda.empty_cache() + return torch.zeros_like(latent_image) + return OwnSampler() + + +def train(model, payload, conds, cfg, adapt, nt, mm, torch, watch, shim): + """Train a LoRA on `model` with ComfyUI's guider carrying each batch. Returns the file's tensors.""" + import comfy.samplers + import comfy_extras.nodes_custom_sampler as custom_sampler + import node_helpers + dtype_name = cfg["training_dtype"] + dtype = node_helpers.string_to_torch_dtype(dtype_name if dtype_name != "none" else "bf16") + lora_dtype = node_helpers.string_to_torch_dtype(cfg["lora_dtype"]) + with torch.inference_mode(False), shim: + model = model.clone(force_deepcopy=True) # the loader made the model's "copy" be the model itself, so this holds one model; made fresh outside inference mode otherwise + model.set_model_compute_dtype(dtype) + model.model.requires_grad_(False).train() + latents, count, multi = nt._prepare_latents_and_count(nt._process_latents_standard_mode(payload), dtype, False) + positive = nt._validate_and_expand_conditioning(nt._process_conditioning(conds), count, False) + diffusion = model.model.diffusion_model + lora = Lora(adapt["lora_targets"](diffusion), cfg["rank"], cfg["rank"], lora_dtype, model.load_device, torch) + jobs.trail("model prepared, starting") + log.info("FunPack LoRA training", f"LoRA on {len(lora.names)} layers, rank {cfg['rank']}, {sum(p.numel() for p in lora.parameters()) / 1e6:.1f}M trainable values") + optimizer = optimizer_for(cfg["optimizer"], lora.parameters(), cfg["learning_rate"], torch) + criterion = nt._create_loss_function(cfg["loss"]) + restore = checkpoint_modules(adapt["checkpoint"](diffusion), torch) if cfg["gradient_checkpointing"] else (lambda: None) + + def stopped(loss): + jobs.step(loss) + mm.throw_exception_if_processing_interrupted() + if jobs.stop_requested(): + from .train import Stopped + raise Stopped() + sampler = Sampler.make(nt, torch, comfy.samplers, custom_sampler, lora, optimizer, criterion, cfg, dtype, stopped) + if cfg["offloading"]: + log.warning("FunPack LoRA training", "ComfyUI's own 'keep weights in system memory' is ignored by this training loop: it has brought ComfyUI down on H3. Use 'Blocks kept in system memory' instead") + guider = nt.TrainGuider(model, offloading=False) + guider.set_conds(positive) + try: + mm.in_training = True + with watch: + nt._run_training_loop(guider, sampler, latents, count, cfg["seed"], False, multi) + finally: + mm.in_training = False + restore() + state = lora.state() if lora else {} + lora.remove() + del optimizer, sampler, guider # a failed run leaves nothing on the card for the next one + return state diff --git a/modules/training/lora/train.py b/modules/training/lora/train.py index 16c19703..a869bd75 100644 --- a/modules/training/lora/train.py +++ b/modules/training/lora/train.py @@ -10,7 +10,7 @@ import os from pathlib import Path from ..._core import log -from . import captions, jobs, media_io, store, swap +from . import captions, jobs, loop, media_io, store, swap POSITIVE = ("own", "liked") NEGATIVE = ("disliked", "bad_composition", "bad_image") @@ -201,18 +201,22 @@ def _run(model, clip, vae, audio_vae, set_name: str, cfg: dict, prompt_id: str = _free_the_card(mm, torch) _card("after freeing", mm, torch) _train_on_a_plain_copy(model) - losses = _loss_hook(nt, cfg["steps"] * cfg["grad_accumulation"]) jobs.say("Training", "training") shim = adapt["shims"]() if adapt and adapt.get("shims") else _nothing() - with shim, losses, _watch(nt, mm, torch, _swapper(adapt, cfg, mm, torch)): - out = nt.TrainLoraNode.execute( - model=[model], latents=payload, positive=conds, batch_size=[batch], steps=[cfg["steps"]], - grad_accumulation_steps=[cfg["grad_accumulation"]], learning_rate=[cfg["learning_rate"]], rank=[cfg["rank"]], - optimizer=[cfg["optimizer"]], loss_function=[cfg["loss"]], seed=[cfg["seed"]], training_dtype=[cfg["training_dtype"]], - lora_dtype=[cfg["lora_dtype"]], quantized_backward=[False], algorithm=[cfg["algorithm"]], - gradient_checkpointing=[cfg["gradient_checkpointing"]], checkpoint_depth=[(adapt or {}).get("checkpoint_depth", 1)], offloading=[cfg["offloading"]], - existing_lora=["[None]"], bucket_mode=[False], bypass_mode=[cfg["bypass"]]) - lora = (out.result if hasattr(out, "result") else out)[0] + watch = _watch(nt, mm, torch, _swapper(adapt, cfg, mm, torch)) + if adapt and adapt.get("lora_targets"): + lora = loop.train(model, payload, conds, cfg, adapt, nt, mm, torch, watch, shim) # our own loop, ComfyUI carries the batches + else: + losses = _loss_hook(nt, cfg["steps"] * cfg["grad_accumulation"]) + with shim, losses, watch: + out = nt.TrainLoraNode.execute( + model=[model], latents=payload, positive=conds, batch_size=[batch], steps=[cfg["steps"]], + grad_accumulation_steps=[cfg["grad_accumulation"]], learning_rate=[cfg["learning_rate"]], rank=[cfg["rank"]], + optimizer=[cfg["optimizer"]], loss_function=[cfg["loss"]], seed=[cfg["seed"]], training_dtype=[cfg["training_dtype"]], + lora_dtype=[cfg["lora_dtype"]], quantized_backward=[False], algorithm=[cfg["algorithm"]], + gradient_checkpointing=[cfg["gradient_checkpointing"]], checkpoint_depth=[(adapt or {}).get("checkpoint_depth", 1)], offloading=[cfg["offloading"]], + existing_lora=["[None]"], bucket_mode=[False], bypass_mode=[cfg["bypass"]]) + lora = (out.result if hasattr(out, "result") else out)[0] file = _save(lora, cfg["name"], folder_paths, safetensors.torch, set_name, cfg) jobs.finish("done", lora=file) return file @@ -309,7 +313,8 @@ def _card(when, mm, torch): device = mm.get_torch_device() gib = 1024 ** 3 loaded = ", ".join(f"{type(getattr(m.model, 'model', m.model)).__name__} {m.model_loaded_memory() / gib:.1f} GiB" for m in mm.current_loaded_models) or "none" - log.info("FunPack LoRA training", f"card {when}: tensors {torch.cuda.memory_allocated(device) / gib:.1f} GiB, free {mm.get_free_memory(device) / gib:.1f} GiB, loaded models: {loaded}") + line = f"card {when}: tensors {torch.cuda.memory_allocated(device) / gib:.1f} GiB, free {mm.get_free_memory(device) / gib:.1f} GiB, loaded models: {loaded}" + jobs.trail(line) except Exception as exc: # noqa: BLE001 -- a report must never stop a run log.info("FunPack LoRA training", f"card {when}: not readable ({type(exc).__name__})") @@ -329,6 +334,8 @@ class _watch: return def watched(*args, **kwargs): + if label == "loading the model": + jobs.trail("loading the model onto the card now") out = original(*args, **kwargs) if label == "loading the model" and self.after_load: self.after_load(args[0] if args else kwargs.get("models"))