move offloading outside of preprocess, which is shared with train

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

Author: kylesayrs

src/llmcompressor/entrypoints/oneshot.py

@@ -2,6 +2,8 @@
 from datetime import datetime
 from typing import Optional
 
+import torch
+from compressed_tensors.utils import force_cpu_offload
 from loguru import logger
 from torch.utils.data import DataLoader
 from transformers import PreTrainedModel
@@ -123,6 +125,15 @@ def __init__(
         # initialize the model and processor
         pre_process(model_args)
 
+        # offload to cpu if possible
+        if "cuda" in str(model_args.oneshot_device) and torch.cuda.is_available():
+            # 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
+            force_cpu_offload(model_args.model, model_args.oneshot_device)
+        else:
+            logger.warning("CUDA is not available! Compressing model on CPU instead")
+
         # Set instance attributes
         self.model = self.model_args.model
         self.processor = self.model_args.processor

src/llmcompressor/entrypoints/utils.py

@@ -3,8 +3,6 @@
 from pathlib import PosixPath
 from typing import Optional, Tuple
 
-import torch
-from compressed_tensors.utils import force_cpu_offload
 from loguru import logger
 from torch.nn import Module
 from transformers import (
@@ -64,16 +62,6 @@ def pre_process(model_args: "ModelArguments"):
     # untie tie_word_embeddings weights
     patch_tied_tensors_bug(model_args.model)
 
-    # offload to cpu if possible
-    if "cuda" in str(model_args.oneshot_device) and torch.cuda.is_available():
-        # TODO: consider renaming function to something like "offload_dispatch_model"
-        # TODO: modify function to remove any hooks if they already exist (making sure
-        # to move to cpu when removing hook
-        force_cpu_offload(model_args.model, model_args.oneshot_device)
-
-    else:
-        logger.warning("CUDA is not available! Compressing model on CPU instead")
-
     # wrap model.save_pretrained
     modify_save_pretrained(model_args.model)