From 6603f42ecbecae1b9a2376d93bd82cbdce89d569 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 17 Oct 2026 20:08:01 +0000 Subject: [PATCH] LoRA training: pause, resume and a preview LoRA at chosen steps (backend) A run can wait at a step boundary (Pause, or every N steps for a preview): its state is saved, the LoRA so far is written as _preview.safetensors, and the card is let go. Resume carries on from the saved state under the same run. UNTESTED on real weights. Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01MpnF7GwNGMbngU11wjym7P --- core/tests/test_shipped_nodes.py | 2 +- modules/training/lora/__init__.py | 9 +++ modules/training/lora/jobs.py | 62 ++++++++++++--- modules/training/lora/loop.py | 80 ++++++++++++++++++-- modules/training/lora/nodes.py | 6 +- modules/training/lora/tests/test_train_h3.py | 46 +++++++++++ modules/training/lora/train.py | 44 +++++++++-- 7 files changed, 223 insertions(+), 26 deletions(-) diff --git a/core/tests/test_shipped_nodes.py b/core/tests/test_shipped_nodes.py index 08939740..0d2c9e43 100644 --- a/core/tests/test_shipped_nodes.py +++ b/core/tests/test_shipped_nodes.py @@ -74,7 +74,7 @@ FROZEN = { "model", "clip", "vae", "audio_vae", "set_name", "name", "use", "steps", "learning_rate", "rank", "algorithm", "optimizer", "loss", "batch_size", "grad_accumulation", "resolution", "aspect_buckets", "frames", "fps", "start", "trigger", "seed", "training_dtype", "lora_dtype", "gradient_checkpointing", "offloading", "bypass", - "model_name", "weight_dtype", "compute_dtype", "postfix", "swap_blocks", + "model_name", "weight_dtype", "compute_dtype", "postfix", "swap_blocks", "preview_every", "resume", ], "FunPackEmptyLatent": ["model", "width", "height", "length", "batch_size", "audio_vae", "frame_rate"], "FunPackDiffusionModelLoader": [ diff --git a/modules/training/lora/__init__.py b/modules/training/lora/__init__.py index 5a640bec..f4b4ee76 100644 --- a/modules/training/lora/__init__.py +++ b/modules/training/lora/__init__.py @@ -204,9 +204,18 @@ def routes(table, base, web): since = 0 return web.json_response(jobs.snapshot(since)) + @table.post(base + "/run/pause") + async def _pause(_req): + return web.json_response({"paused": jobs.request_pause()}) + @table.post(base + "/run/stop") async def _stop(_req): + now = jobs.snapshot() + was_paused = now["state"] == "paused" stopped = jobs.request_stop() + if stopped and was_paused: + train.discard(now["run"], now["name"]) + return web.json_response({"stopped": True}) if stopped: try: import comfy.model_management diff --git a/modules/training/lora/jobs.py b/modules/training/lora/jobs.py index b60d7baf..3c991438 100644 --- a/modules/training/lora/jobs.py +++ b/modules/training/lora/jobs.py @@ -19,10 +19,11 @@ MAX_KEPT_LOSSES = 4000 # a saved run keeps a thinned curve: a 100000 st _LOCK = threading.Lock() _ID = re.compile(r"\A[0-9a-f]{12}\Z") IDLE = {"state": "idle", "run": "", "set": "", "name": "", "step": 0, "total": 0, "started": 0.0, "message": "", "error": "", - "lora": "", "samples": [], "settings": {}} + "lora": "", "samples": [], "settings": {}, "paused": None} _now: dict = dict(IDLE) _loss: list[float] = [] _stop = False +_pause = False def _dir(): @@ -42,14 +43,19 @@ def trail(message: str) -> None: 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 +def begin(set_name: str, name: str, total: int, settings: dict, prompt_id: str = "", resume: dict | None = None) -> str: + """Start a record. A run that was still marked running (ComfyUI restarted under it) is simply replaced. + + `resume` ({"run", "loss"}) carries on a paused run: same record, same loss curve, the steps counting on. + """ + global _stop, _pause with _LOCK: - _stop = False + _stop = _pause = False _loss.clear() _now.clear() - _now.update(IDLE, state="training", run=uuid.uuid4().hex[:12], set=set_name, name=name, total=int(total), started=time.time(), + if resume: + _loss.extend(float(v) for v in resume.get("loss") or []) + _now.update(IDLE, state="training", run=(resume or {}).get("run") or uuid.uuid4().hex[:12], step=len(_loss), set=set_name, name=name, total=int(total), started=time.time(), settings=dict(settings), prompt_id=prompt_id, message="Loading the model and the training data") return _now["run"] @@ -84,12 +90,40 @@ def thin(values, keep: int = MAX_KEPT_LOSSES) -> list[float]: return out + [values[-1]] +def losses() -> list[float]: + with _LOCK: + return list(_loss) + + +def pause(reason: str, step: int, preview: str, checkpoint: bool = True) -> None: + """The run waits: its state is on disk, and the models are let go so the pipeline can use the card. `reason`: user or preview.""" + with _LOCK: + if _now["state"] != "training": + return + _now.update(state="paused", message="", paused={"reason": reason, "step": int(step), "preview": preview, "checkpoint": bool(checkpoint)}) + + +def request_pause() -> bool: + """Ask the running training to wait at its next step. False when nothing is training.""" + global _pause + with _LOCK: + if _now["state"] != "training": + return False + _pause = True + _now["message"] = "Pausing" + return True + + +def pause_requested() -> bool: + return _pause + + def finish(state: str, lora: str = "", error: str = "") -> None: """state: done, stopped (the person stopped it) or error. Writes the run's record.""" with _LOCK: if _now["run"] == "" or _now["state"] in ("done", "error", "stopped"): return - _now.update(state=state, lora=lora, error=error, message="") + _now.update(state=state, lora=lora, error=error, message="", paused=None) record = {k: _now.get(k) for k in ("run", "set", "name", "total", "started", "settings", "lora", "samples", "error")} record.update(state=state, steps=len(_loss), seconds=round(time.time() - _now["started"], 1), final_loss=_loss[-1] if _loss else None, loss=thin(_loss)) @@ -144,11 +178,17 @@ def request_stop() -> bool: """Ask the running training to end at its next step. False when nothing is training.""" global _stop with _LOCK: - if _now["state"] != "training": + if _now["state"] == "paused": + ended = True # nothing is running: the run simply ends + elif _now["state"] != "training": return False - _stop = True - _now["message"] = "Stopping" - return True + else: + ended = False + _stop = True + _now["message"] = "Stopping" + if ended: + finish("stopped") + return True def stop_requested() -> bool: diff --git a/modules/training/lora/loop.py b/modules/training/lora/loop.py index ff2e9c60..273ed87b 100644 --- a/modules/training/lora/loop.py +++ b/modules/training/lora/loop.py @@ -15,6 +15,45 @@ from ..._core import log from . import jobs +class Paused(Exception): + """The loop stops at a step boundary and keeps its state: `reason` is user (Pause pressed) or preview (the step asked for).""" + + def __init__(self, reason, step): + super().__init__(f"paused at step {step}") + self.reason, self.step = reason, step + + +def _checkpoint_file(run): + from ..._core import config + folder = config.TRAINING_DIR / "checkpoints" + folder.mkdir(parents=True, exist_ok=True) + return folder / f"{run}.pt" + + +def save_checkpoint(run, lora, optimizer, step, torch): + """Everything a resumed run needs besides the data: the LoRA, the optimizer's memory, the loss so far and where it stopped.""" + torch.save({"lora": {n: tuple(p.detach().cpu() for p in pair) for n, pair in lora.pairs.items()}, "optimizer": optimizer.state_dict(), + "step": int(step), "loss": jobs.losses()}, _checkpoint_file(run)) + + +def load_checkpoint(run, torch): + """The saved state of a paused run, or a ValueError saying why it cannot be carried on.""" + import re + if not re.fullmatch(r"[0-9a-f]{12}", str(run)): + raise ValueError("that is not the name of a paused run") + path = _checkpoint_file(run) + if not path.is_file(): + raise ValueError("the paused run's saved state is gone: start the training again") + return torch.load(path, map_location="cpu", weights_only=False) + + +def forget_checkpoint(run): + try: + _checkpoint_file(run).unlink() + except OSError: + pass + + 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.""" @@ -41,6 +80,17 @@ class Lora: return out + delta.to(out.dtype) return add + def restore(self, saved): + """Put a checkpoint's weights back; the layers must be the ones it was saved from.""" + if set(saved) != set(self.pairs): + raise ValueError("the paused run was made for a different model or LoRA layout: start the training again") + with self.torch.no_grad(): + for name, (down, up) in self.pairs.items(): + if saved[name][0].shape != down.shape: + raise ValueError("the paused run used another rank: start the training again") + down.copy_(saved[name][0]) + up.copy_(saved[name][1]) + def parameters(self): return [p for pair in self.pairs.values() for p in pair] @@ -88,7 +138,7 @@ 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): + def make(nt, torch, comfy_samplers, custom_sampler, lora, optimizer, criterion, cfg, dtype, stopped, first=0): 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) @@ -96,7 +146,9 @@ class Sampler: 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): + total = cfg["steps"] * accumulate + every = -(-int(cfg.get("preview_every") or 0) // accumulate) * accumulate # a preview waits for a whole optimizer step + for i in range(first, total): 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]) @@ -121,12 +173,17 @@ class Sampler: if (i + 1) % accumulate == 0: optimizer.step() optimizer.zero_grad() + if (i + 1) < total: + if jobs.pause_requested(): + raise Paused("user", i + 1) + if every and (i + 1) % every == 0: + raise Paused("preview", i + 1) torch.cuda.empty_cache() return torch.zeros_like(latent_image) return OwnSampler() -def train(model, payload, conds, cfg, adapt, nt, mm, torch, watch, shim): +def train(model, payload, conds, cfg, adapt, nt, mm, torch, watch, shim, resume=None, run=None): """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 @@ -154,19 +211,30 @@ def train(model, payload, conds, cfg, adapt, nt, mm, torch, watch, shim): 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) + first = 0 + if resume: + lora.restore(resume["lora"]) + optimizer.load_state_dict(resume["optimizer"]) + first = int(resume["step"]) + log.info("FunPack LoRA training", f"carrying on from step {first}") + sampler = Sampler.make(nt, torch, comfy.samplers, custom_sampler, lora, optimizer, criterion, cfg, dtype, stopped, first) 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) + paused = None try: mm.in_training = True with watch: - nt._run_training_loop(guider, sampler, latents, count, cfg["seed"], False, multi) + try: + nt._run_training_loop(guider, sampler, latents, count, cfg["seed"], False, multi) + except Paused as stop: + paused = stop + save_checkpoint(run, lora, optimizer, stop.step, torch) 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 + return state, paused diff --git a/modules/training/lora/nodes.py b/modules/training/lora/nodes.py index 7e711081..1c1bedf4 100644 --- a/modules/training/lora/nodes.py +++ b/modules/training/lora/nodes.py @@ -49,6 +49,8 @@ class FunPackTrainLora(io.ComfyNode): io.Combo.Input("compute_dtype", options=["default", "fp16", "bf16", "fp32"], default="default", tooltip="What the maths runs in when the model is loaded here."), io.String.Input("postfix", default="", tooltip="Text put at the end of every caption."), io.Int.Input("swap_blocks", default=0, min=0, max=128, tooltip="How many transformer blocks keep their weights in system memory, fetched one layer at a time: slower, and each one frees about a block's weight on the card. 0 keeps the whole model on the card."), + io.Int.Input("preview_every", default=0, min=0, max=100000, tooltip="Pause every this many steps so the pipeline can make a preview with the LoRA so far. 0 never pauses."), + io.String.Input("resume", default="", tooltip="The name of a paused run to carry on. Empty starts a new run."), ], outputs=[io.String.Output(display_name="lora")], ) @@ -56,7 +58,7 @@ class FunPackTrainLora(io.ComfyNode): @classmethod def execute(cls, clip, vae, set_name, name, use, steps, learning_rate, rank, algorithm, optimizer, loss, batch_size, grad_accumulation, resolution, aspect_buckets, frames, fps, start, trigger, seed, training_dtype, lora_dtype, - gradient_checkpointing, offloading, bypass, model=None, audio_vae=None, model_name="", weight_dtype="default", compute_dtype="default", postfix="", swap_blocks=0) -> io.NodeOutput: + gradient_checkpointing, offloading, bypass, model=None, audio_vae=None, model_name="", weight_dtype="default", compute_dtype="default", postfix="", swap_blocks=0, preview_every=0, resume="") -> io.NodeOutput: if not addons.enabled("lora_training"): raise RuntimeError("LoRA training is switched off. Turn it on in Settings > Modules.") cfg = {"name": name, "use": use, "steps": steps, "learning_rate": learning_rate, "rank": rank, "algorithm": algorithm, @@ -64,5 +66,5 @@ class FunPackTrainLora(io.ComfyNode): "resolution": resolution, "aspect_buckets": aspect_buckets, "frames": frames, "fps": fps, "start": start, "trigger": trigger, "seed": seed, "training_dtype": training_dtype, "lora_dtype": lora_dtype, "gradient_checkpointing": gradient_checkpointing, "offloading": offloading, "bypass": bypass, - "model_name": model_name, "weight_dtype": weight_dtype, "compute_dtype": compute_dtype, "postfix": postfix, "swap_blocks": swap_blocks} + "model_name": model_name, "weight_dtype": weight_dtype, "compute_dtype": compute_dtype, "postfix": postfix, "swap_blocks": swap_blocks, "preview_every": preview_every, "resume": resume} return io.NodeOutput(train.run(model, clip, vae, audio_vae, set_name, cfg)) diff --git a/modules/training/lora/tests/test_train_h3.py b/modules/training/lora/tests/test_train_h3.py index 07b99594..a4b180c5 100644 --- a/modules/training/lora/tests/test_train_h3.py +++ b/modules/training/lora/tests/test_train_h3.py @@ -351,3 +351,49 @@ def test_block_modulation_is_out_of_place_while_training_and_put_back_after(): assert torch.allclose(m._mod_gate(x0, gate, other, segs), want_x) assert torch.equal(x0, x) assert (m._mod_scale_shift, m._mod_gate) == before + + +def test_a_run_pauses_for_a_preview_and_carries_on_to_the_same_end(run_env): + import safetensors.torch + from modules.training.lora import train + tmp_path, jobs = run_env + first = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2)) + snap = jobs.snapshot() + assert first == "" and snap["state"] == "paused" + assert snap["paused"]["reason"] == "preview" and snap["paused"]["step"] == 2 and snap["step"] == 2 + preview = tmp_path / "loras" / snap["paused"]["preview"] + assert preview.name == "tiny_lora_preview.safetensors" + assert any(k.endswith(".lora_up.weight") for k in safetensors.torch.load_file(str(preview))) + run = snap["run"] + # the person looks at the preview, and training carries on from where it stopped + second = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2, resume=run)) + snap = jobs.snapshot() + assert second == "" and snap["run"] == run and snap["paused"]["step"] == 4 and len(snap["loss"]) == 4 + done = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2, resume=run)) + snap = jobs.snapshot() + assert done == "tiny_lora.safetensors" and snap["state"] == "done" and len(snap["loss"]) == 6 + assert not preview.exists(), "the preview file stays after the run is over" + assert not (tmp_path / "training" / "checkpoints" / f"{run}.pt").exists() + + +def test_pause_button_waits_at_the_next_step_and_stop_ends_a_paused_run(run_env, monkeypatch): + from modules.training.lora import jobs as j, train + tmp_path, jobs = run_env + real = j.step + + def press(loss): + real(loss) + j.request_pause() + monkeypatch.setattr(j, "step", press) + train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6)) + snap = jobs.snapshot() + assert snap["state"] == "paused" and snap["paused"]["reason"] == "user" and snap["paused"]["step"] == 1 + assert j.request_stop() and jobs.snapshot()["state"] == "stopped" + + +def test_a_run_that_cannot_be_resumed_says_why(run_env): + from modules.training.lora import train + _, jobs = run_env + with pytest.raises(ValueError, match="saved state is gone"): + train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(resume="0123456789ab")) + assert jobs.snapshot()["state"] == "error" diff --git a/modules/training/lora/train.py b/modules/training/lora/train.py index a869bd75..55ed0ca4 100644 --- a/modules/training/lora/train.py +++ b/modules/training/lora/train.py @@ -128,7 +128,10 @@ def run(model, clip, vae, audio_vae, set_name: str, cfg: dict, prompt_id: str = torch.cuda.empty_cache() _card("at the start of the run", mm, torch) try: - return _run(model, clip, vae, audio_vae, set_name, cfg, prompt_id) + result = _run(model, clip, vae, audio_vae, set_name, cfg, prompt_id) + model = clip = vae = audio_vae = None # a paused or finished run lets the card go: the pipeline's preview or the next run needs it + _free_the_card(mm, torch) + return result except BaseException as exc: traceback.clear_frames(exc.__traceback__) model = clip = vae = audio_vae = None @@ -144,8 +147,16 @@ def _run(model, clip, vae, audio_vae, set_name: str, cfg: dict, prompt_id: str = import safetensors.torch from ..._core import registry - run_id = jobs.begin(set_name, cfg["name"], cfg["steps"], cfg, prompt_id) # first, so a refusal below is shown in the window and not only in ComfyUI's log + resume, failed = None, None + if cfg.get("resume"): + try: + resume = loop.load_checkpoint(cfg["resume"], torch) + except ValueError as exc: + failed = exc + run_id = jobs.begin(set_name, cfg["name"], cfg["steps"], cfg, prompt_id, {"run": cfg["resume"], "loss": resume["loss"]} if resume else None) # first, so a refusal below is shown in the window and not only in ComfyUI's log try: + if failed: + raise failed if model is None: model = load_model(cfg.get("model_name"), cfg.get("weight_dtype", "default"), cfg.get("compute_dtype", "default")) adapt = registry.current().ask("training_adapt", model) @@ -205,7 +216,12 @@ def _run(model, clip, vae, audio_vae, set_name: str, cfg: dict, prompt_id: str = shim = adapt["shims"]() if adapt and adapt.get("shims") else _nothing() 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 + lora, paused = loop.train(model, payload, conds, cfg, adapt, nt, mm, torch, watch, shim, resume, run_id) # our own loop, ComfyUI carries the batches + if paused: + preview = _save(lora, f"{cfg['name']}_preview", folder_paths, safetensors.torch, set_name, cfg, keep=True) + jobs.pause(paused.reason, paused.step, preview) + log.info("FunPack LoRA training", f"paused at step {paused.step} ({paused.reason}); the LoRA so far is {preview}") + return "" else: losses = _loss_hook(nt, cfg["steps"] * cfg["grad_accumulation"]) with shim, losses, watch: @@ -218,6 +234,7 @@ def _run(model, clip, vae, audio_vae, set_name: str, cfg: dict, prompt_id: str = 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) + discard(run_id, cfg["name"]) jobs.finish("done", lora=file) return file except Stopped: @@ -467,13 +484,28 @@ class _loss_hook: return False -def _save(lora, name, folder_paths, safetensors_torch, set_name, cfg) -> str: - """Into the first loras folder, as .safetensors, or _2 and so on rather than over an earlier one.""" +def _preview_name(name) -> str: + stem = "".join(c for c in f"{name}_preview" if c.isalnum() or c in "-_ .").strip() or "funpack_lora_preview" + return f"{stem}.safetensors" + + +def discard(run, name) -> None: + """A run that is over (finished, stopped) keeps neither its saved state nor its preview file.""" + import folder_paths + loop.forget_checkpoint(run) + try: + (Path(folder_paths.get_folder_paths("loras")[0]) / _preview_name(name.removesuffix("_preview"))).unlink() + except OSError: + pass + + +def _save(lora, name, folder_paths, safetensors_torch, set_name, cfg, keep=False) -> str: + """Into the first loras folder, as .safetensors, or _2 and so on rather than over an earlier one. `keep`: replace the file of that name (the preview).""" folder = Path(folder_paths.get_folder_paths("loras")[0]) folder.mkdir(parents=True, exist_ok=True) stem = "".join(c for c in name if c.isalnum() or c in "-_ .") .strip() or "funpack_lora" n, path = 1, folder / f"{stem}.safetensors" - while path.exists(): + while path.exists() and not keep: n += 1 path = folder / f"{stem}_{n}.safetensors" meta = {"software": "ComfyUI-FunPack", "training_set": set_name, "algorithm": cfg["algorithm"], "rank": str(cfg["rank"]),