From 4ed6ae8726aa8545e251447bb9c0024335e3a7ae Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 17 Nov 2024 16:54:44 -0800 Subject: [PATCH] Update _utils.py --- unsloth/models/_utils.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index adcab5e980..e07f5ac2ce 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -55,6 +55,10 @@ __all__ = [ "HAS_CUT_CROSS_ENTROPY", "fused_linear_cross_entropy", + "patch_unsloth_smart_gradient_checkpointing", + "unpatch_unsloth_smart_gradient_checkpointing", + + "patch_compiled_autograd", ] import torch @@ -73,6 +77,7 @@ from unsloth_zoo.patching_utils import ( patch_layernorm, patch_torch_compile, patch_model_and_tokenizer, + patch_compiled_autograd, ) from unsloth_zoo.gradient_checkpointing import ( Unsloth_Offloaded_Gradient_Checkpointer, @@ -84,6 +89,9 @@ from unsloth_zoo.gradient_checkpointing import ( unsloth_gradient_checkpoint, patch_gradient_checkpointing, unpatch_gradient_checkpointing, + + patch_unsloth_smart_gradient_checkpointing, + unpatch_unsloth_smart_gradient_checkpointing, ) from unsloth_zoo.loss_utils import ( HAS_CUT_CROSS_ENTROPY,