From b0afaab8dbbf40f93e98caad40585df951954af7 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 17 Oct 2026 18:33:26 +0000 Subject: [PATCH] Training loads the model and copies its inputs outside inference mode Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01MpnF7GwNGMbngU11wjym7P --- modules/training/lora/tests/test_train_h3.py | 30 +++++++++++++++++++- modules/training/lora/train.py | 19 ++++++++++++- 2 files changed, 47 insertions(+), 2 deletions(-) diff --git a/modules/training/lora/tests/test_train_h3.py b/modules/training/lora/tests/test_train_h3.py index 5c0760aa..a9d0da6d 100644 --- a/modules/training/lora/tests/test_train_h3.py +++ b/modules/training/lora/tests/test_train_h3.py @@ -246,6 +246,9 @@ def test_the_trainer_loads_the_model_file_once_and_its_copy_is_that_model(monkey seen = {} class Patcher: + def __init__(self): + self.weight = torch.zeros(2) # made where the loader runs: inference mode must not reach it + def set_model_compute_dtype(self, dtype): seen["compute"] = dtype @@ -254,7 +257,9 @@ def test_the_trainer_loads_the_model_file_once_and_its_copy_is_that_model(monkey return Patcher() monkeypatch.setattr(comfy.sd, "load_diffusion_model", load) monkeypatch.setattr(folder_paths, "get_full_path_or_raise", lambda kind, name: f"/models/{name}") - model = train.load_model("m.safetensors", "bf16", "fp32") + with torch.inference_mode(): + model = train.load_model("m.safetensors", "bf16", "fp32") + assert not model.weight.is_inference() assert seen["path"] == "/models/m.safetensors" and seen["options"] == {"dtype": torch.bfloat16} and seen["dynamic"] and seen["compute"] == torch.float32 factory, args = model.cached_patcher_init assert factory(*args, disable_dynamic=True) is model and seen["loads"] == 1 # the trainer's copy reads nothing again @@ -273,3 +278,26 @@ def test_h3_is_checkpointed_one_block_at_a_time_not_as_a_whole(run_env): pieces = nt.find_modules_at_depth(net, depth=depth) assert depth == 2 and len(pieces) > len(net.blocks) and all(b in pieces for b in net.blocks) assert net not in pieces + + +def test_what_was_made_in_inference_mode_is_copied_for_the_backward_pass(): + import torch + from modules.training.lora import train + with torch.inference_mode(): + made = [(torch.ones(2), {"k": torch.zeros(1), "n": 3})], [torch.ones(1)] + assert made[1][0].is_inference() + got = train._normal(made) + assert not got[1][0].is_inference() and not got[0][0][0].is_inference() and not got[0][0][1]["k"].is_inference() and got[0][0][1]["n"] == 3 + x = got[1][0].requires_grad_(False) + (x * torch.ones(1, requires_grad=True)).sum().backward() # a saved-for-backward use no longer fails + + +def test_a_tiny_h3_trains_when_run_the_way_comfyui_runs_a_node(run_env): + """ComfyUI runs every node in inference mode, and the models, text and latents here are made inside it.""" + import torch + from modules.training.lora import train + _, jobs = run_env + with torch.inference_mode(): + patcher = tiny_h3_patcher() + train.run(patcher, FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=2)) + assert jobs.snapshot()["state"] == "done" diff --git a/modules/training/lora/train.py b/modules/training/lora/train.py index b6f34fe4..276d1f2b 100644 --- a/modules/training/lora/train.py +++ b/modules/training/lora/train.py @@ -163,6 +163,7 @@ def run(model, clip, vae, audio_vae, set_name: str, cfg: dict, prompt_id: str = videos.append(latent.clone()) sounds.append(None if sound is None else sound.clone()) conds.append(_encode_text(clip, text)) + videos, sounds, conds = _normal(videos), _normal(sounds), _normal(conds) # made in inference mode: copies a backward pass can use if adapt and adapt.get("pack") and audio_vae is not None and all(s is not None for s in sounds): latents = adapt["pack"](videos, sounds) payload = [latents] @@ -229,13 +230,29 @@ def load_model(name, weight_dtype="default", compute_dtype="default"): options = {} if weight_dtype in _WEIGHT: options["dtype"] = getattr(torch, _WEIGHT[weight_dtype]) - model = comfy.sd.load_diffusion_model(path, model_options=options, disable_dynamic=True) + with torch.inference_mode(False): # ComfyUI runs nodes in inference mode; weights made there cannot take part in a backward pass + model = comfy.sd.load_diffusion_model(path, model_options=options, disable_dynamic=True) if compute_dtype in _WEIGHT and hasattr(model, "set_model_compute_dtype"): model.set_model_compute_dtype(getattr(torch, _WEIGHT[compute_dtype])) model.cached_patcher_init = (lambda *args, **kwargs: model, ()) # the trainer's "fresh copy" is this model: no second read of the file return model +def _normal(value): + """Tensors made in inference mode, copied so autograd may save them; lists, tuples and dicts are walked.""" + import torch + if isinstance(value, torch.Tensor): + with torch.inference_mode(False): + return value.clone() if value.is_inference() else value + if isinstance(value, list): + return [_normal(v) for v in value] + if isinstance(value, tuple): + return tuple(_normal(v) for v in value) + if isinstance(value, dict): + return {k: _normal(v) for k, v in value.items()} + return value + + def _fit_frames(cfg, adapt, items): """Every clip is read for the same number of frames, so the shortest clip sets the limit: fewer than asked is used, and said.""" rate = cfg["fps"] or adapt.get("fps", 24)