Training loads the model and copies its inputs outside inference mode

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MpnF7GwNGMbngU11wjym7P

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"

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)