remove hooks

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

src/llmcompressor/entrypoints/oneshot.py

@@ -14,6 +14,8 @@
 from llmcompressor.entrypoints.utils import post_process, pre_process
 from llmcompressor.pipelines import CalibrationPipeline
 
+from accelerate.hooks import remove_hook_from_module
+
 __all__ = ["Oneshot", "oneshot"]
 
 
@@ -130,6 +132,8 @@ def __init__(
             # TODO: consider renaming function similar to "offload_dispatch_model"
             # TODO: modify function to remove any hooks if they already exist (making
             # sure to move to cpu when removing hook
+            # TODO: remove hook in util
+            remove_hook_from_module(model_args.model, recurse=True)
             force_cpu_offload(model_args.model, model_args.oneshot_device)
         else:
             logger.warning("CUDA is not available! Compressing model on CPU instead")

src/llmcompressor/entrypoints/utils.py

@@ -27,6 +27,7 @@
 )
 from llmcompressor.typing import Processor
 from llmcompressor.utils.fsdp.helpers import is_fsdp_model
+from accelerate.hooks import remove_hook_from_module
 
 
 def pre_process(model_args: "ModelArguments"):
@@ -105,6 +106,9 @@ def post_process(
             "Ex. `oneshot(..., output_dir=...)`"
         )
 
+    # Remove any existing hooks (maybe added by oneshot sequential onloading)
+    remove_hook_from_module(model_args.model, recurse=True)
+
     # Reset the one-time-use session upon completion
     if recipe_args is not None and recipe_args.clear_sparse_session:
         reset_session()