Add skip_missing_weights_context

src/llmcompressor/transformers/sparsification/compressed_tensors_utils.py

@@ -1,3 +1,5 @@
+import contextlib
+import logging
 import os
 import re
 import weakref
@@ -148,15 +150,6 @@ def save_pretrained_wrapper(
             # https://github.com/huggingface/transformers/pull/30488
             transformers.modeling_utils.dtype_byte_size = new_dtype_byte_size
 
-            def skip(*args, **kwargs):
-                pass
-
-            # Skip the initializer step. This accelerates the loading
-            # of the models, especially for the quantized models
-            torch.nn.init.kaiming_uniform_ = skip
-            torch.nn.init.uniform_ = skip
-            torch.nn.init.normal_ = skip
-
             # state_dict gets passed in as a kwarg for FSDP models
             state_dict = kwargs.pop("state_dict", None)
             if state_dict is None:
@@ -305,3 +298,40 @@ def get_model_compressor(
         sparsity_config=sparsity_config,
         quantization_format=quantization_format,
     )
+
+
+@contextlib.contextmanager
+def skip_missing_weights_context():
+    """
+    Used when loading a quantized model whose state dict does not align model
+    definition weights
+    """
+    kaiming_uniform_ = torch.nn.init.kaiming_uniform_
+    uniform_ = torch.nn.init.uniform_
+    normal_ = torch.nn.init.normal_
+
+    transformers_logger = logging.getLogger("transformers.modeling_utils")
+    restore_log_level = transformers_logger.getEffectiveLevel()
+
+    # skip init functions
+    def skip(*args, **kwargs):
+        pass
+
+    torch.nn.init.kaiming_uniform_ = skip
+    torch.nn.init.uniform_ = skip
+    torch.nn.init.normal_ = skip
+    # TODO: consider skipping other default init functions
+
+    # temporarily set the log level to error, to ignore printing out long missing
+    # and unexpected key error messages (these are EXPECTED for quantized models)
+    transformers_logger.setLevel(level=logging.ERROR)
+
+    yield
+
+    # restore original functions
+    torch.nn.init.kaiming_uniform_ = kaiming_uniform_
+    torch.nn.init.uniform_ = uniform_
+    torch.nn.init.normal_ = normal_
+
+    # restore transformers logging level now that model shell is loaded
+    transformers_logger.setLevel(level=restore_log_level)

tests/llmcompressor/transformers/sparsification/test_compress_tensor_utils.py

@@ -1,4 +1,3 @@
-import logging
 import math
 import shutil
 
@@ -22,6 +21,7 @@
 from llmcompressor.transformers.sparsification.compressed_tensors_utils import (
     modify_save_pretrained,
     patch_tied_tensors_bug,
+    skip_missing_weights_context,
 )
 
 
@@ -63,18 +63,10 @@ def test_sparse_model_reload(compressed, config, dtype, tmp_path):
         clear_sparse_session=False,
     )
 
-    # temporarily set the log level to error, to ignore printing out long missing
-    # and unexpected key error messages (these are EXPECTED for quantized models)
-    transformers_logger = logging.getLogger("transformers.modeling_utils")
-    restore_log_level = transformers_logger.getEffectiveLevel()
-    transformers_logger.setLevel(level=logging.ERROR)
-
-    model = AutoModelForCausalLM.from_pretrained(
-        tmp_path / "oneshot_out", torch_dtype=dtype
-    )
-
-    # restore transformers logging level now that model shell is loaded
-    transformers_logger.setLevel(level=restore_log_level)
+    with skip_missing_weights_context():
+        model = AutoModelForCausalLM.from_pretrained(
+            tmp_path / "oneshot_out", torch_dtype=dtype
+        )
 
     # assert that sample layer has the intended sparsity
     assert math.isclose(