[Offload] Convert model back to CT offloading for testing (vllm-project#2403)

## Purpose ##
* Fix testing which requires access to the model after the model has
been saved

## Prerequisites ##
* vllm-project/compressed-tensors#601

## Changes ##
* Convert back to CT offloading after converting to accelerate
offloading for saving
* Previously we just "removed dispatch", but this is bad practice as it
won't work for disk offloading

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Co-authored-by: Dipika Sikka <dipikasikka1@gmail.com>

Author: kylesayrs

src/llmcompressor/transformers/compression/compressed_tensors_utils.py

@@ -9,7 +9,7 @@
     SparsityCompressionConfig,
 )
 from compressed_tensors.config import CompressionFormat
-from compressed_tensors.offload import is_rank0, to_accelerate
+from compressed_tensors.offload import from_accelerate, is_rank0, to_accelerate
 from loguru import logger
 from transformers import PreTrainedModel
 
@@ -89,10 +89,10 @@ def save_pretrained_wrapper(
             if compressor is not None:
                 compressor.compress_model(model)
 
-            if is_rank0():
-                # convert to accelerate offloaded for optimal saving with transformers
-                to_accelerate(model)
+            # convert to accelerate offloaded for optimal saving with transformers
+            to_accelerate(model)
 
+            if is_rank0():
                 # save (compressed) model structure
                 original_save_pretrained.__get__(model, model_class)(
                     save_directory,
@@ -110,6 +110,9 @@ def save_pretrained_wrapper(
                 # copy python files from cache dir to save_path if any
                 copy_python_files_from_model_cache(model, save_directory)
 
+            # convert back from accelerate to restore model to original form
+            from_accelerate(model)
+
         save_pretrained_wrapper._overridden = True
         return save_pretrained_wrapper