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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MpnF7GwNGMbngU11wjym7P

.gitignore

@@ -49,3 +49,4 @@ app/playwright-report/
 /block_influence.latest.json
 /phrase_probe.enabled
 /phrase_probe.latest.json
+/log.jsonl

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"

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()

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"]

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

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"):

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

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"))