diff --git a/modules/models/minimax_h3/training.py b/modules/models/minimax_h3/training.py index 5af3e9d3..de623469 100644 --- a/modules/models/minimax_h3/training.py +++ b/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} diff --git a/modules/training/lora/tests/test_train_h3.py b/modules/training/lora/tests/test_train_h3.py index 4df4fa7f..5c0760aa 100644 --- a/modules/training/lora/tests/test_train_h3.py +++ b/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 diff --git a/modules/training/lora/train.py b/modules/training/lora/train.py index 3c6ba2cc..1d20c711 100644 --- a/modules/training/lora/train.py +++ b/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)