Enable gradient checkpointing

2025-05-17 20:40:21 +08:00
parent 330cc0eeb4
commit 3d6fe064f5
3 changed files with 5 additions and 20 deletions
@@ -68,6 +68,7 @@ class Pipeline(LightningModule):
        acestep_pipeline.load_checkpoint(acestep_pipeline.checkpoint_dir)

        transformers = acestep_pipeline.ace_step_transformer.float().cpu()
+        transformers.enable_gradient_checkpointing()

        assert lora_config_path is not None, "Please provide a LoRA config path"
        if lora_config_path is not None: