H3 training checkpoints one block at a time and the text refiner on its own

Depth 1 made the whole model a single checkpoint, so nothing was saved.

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

modules/models/minimax_h3/training.py

@@ -71,5 +71,6 @@ def adapt(model):
     if not is_h3(model):
         return None
     return {"name": "H3", "step": SPATIAL_STEP, "video": True, "batch_limit": 1, "fps": FPS,
+            "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))),
             "temporal": temporal_shape, "shims": shims, "pack": pack, "silence": silence_latent}

modules/training/lora/tests/test_train_h3.py

@@ -262,3 +262,14 @@ def test_the_trainer_loads_the_model_file_once_and_its_copy_is_that_model(monkey
         train.load_model("m.gguf")
     with pytest.raises(ValueError, match="no model"):
         train.load_model("")
+
+
+def test_h3_is_checkpointed_one_block_at_a_time_not_as_a_whole(run_env):
+    import comfy_extras.nodes_train as nt
+    from modules.models.minimax_h3 import training
+    patcher = tiny_h3_patcher()
+    depth = training.adapt(patcher)["checkpoint_depth"]
+    net = patcher.model.diffusion_model
+    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

modules/training/lora/train.py

@@ -186,7 +186,7 @@ def run(model, clip, vae, audio_vae, set_name: str, cfg: dict, prompt_id: str =
                 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=[1], offloading=[cfg["offloading"]],
+                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)