revert submodule to module

Author: horheynm

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -42,7 +42,7 @@
 from compressed_tensors.quantization.lifecycle import expand_sparse_target_names
 from compressed_tensors.quantization.quant_args import QuantizationArgs
 from compressed_tensors.quantization.utils import (
-    is_submodule_quantized,
+    is_model_quantized,
     iter_named_leaf_modules,
 )
 from compressed_tensors.utils import get_safetensors_folder, update_parameter_data
@@ -426,7 +426,7 @@ def map_modules_to_quant_args(model: Module) -> Dict[str, QuantizationArgs]:
     """
     quantized_modules_to_args = {}
     for name, submodule in iter_named_leaf_modules(model):
-        if is_submodule_quantized(submodule):
+        if is_model_quantized(submodule):
             if submodule.quantization_scheme.weights is not None:
                 name = fix_fsdp_module_name(name)
                 quantized_modules_to_args[name] = submodule.quantization_scheme.weights

src/compressed_tensors/quantization/lifecycle/apply.py

@@ -56,7 +56,7 @@
     "is_sparse_target",
 ]
 
-from compressed_tensors.quantization.utils.helpers import is_submodule_quantized
+from compressed_tensors.quantization.utils.helpers import is_model_quantized
 from compressed_tensors.utils.safetensors_load import get_quantization_state_dict
 
 
@@ -76,7 +76,7 @@ def load_pretrained_quantization(model: Module, model_name_or_path: str):
     state_dict = get_quantization_state_dict(model_path)
 
     for name, submodule in iter_named_leaf_modules(model):
-        if not is_submodule_quantized(submodule):
+        if not is_model_quantized(submodule):
             continue
         if submodule.quantization_scheme.weights is not None:
             base_name = "weight"

src/compressed_tensors/quantization/quant_config.py

@@ -23,7 +23,7 @@
 )
 from compressed_tensors.quantization.utils import (
     calculate_compression_ratio,
-    is_submodule_quantized,
+    is_model_quantized,
     iter_named_quantizable_modules,
     module_type,
     parse_out_kv_cache_args,
@@ -181,7 +181,7 @@ def from_pretrained(
             model, include_children=True, include_attn=True
         ):
             layer_type = module_type(submodule)
-            if not is_submodule_quantized(submodule):
+            if not is_model_quantized(submodule):
                 if layer_type not in ignore:
                     ignore[layer_type] = []
                 ignore[layer_type].append(name)

src/compressed_tensors/quantization/utils/helpers.py

@@ -41,7 +41,7 @@
     "is_kv_cache_quant_scheme",
     "is_model_quantized",
     "is_model_quantized_from_path",
-    "is_submodule_quantized",
+    "is_model_quantized",
     "iter_named_leaf_modules",
     "iter_named_quantizable_modules",
     "module_type",
@@ -169,7 +169,7 @@ def infer_quantization_status(model: Module) -> Optional["QuantizationStatus"]:
     return None
 
 
-def is_submodule_quantized(module: Module) -> bool:
+def is_model_quantized(module: Module) -> bool:
     """
     Check if a module is quantized, based on the existence of a non-empty quantization
     scheme
@@ -202,7 +202,7 @@ def is_model_quantized(model: Module) -> bool:
     """
 
     for _, submodule in iter_named_leaf_modules(model):
-        if is_submodule_quantized(submodule):
+        if is_model_quantized(submodule):
             return True
 
     return False
@@ -353,7 +353,7 @@ def calculate_compression_ratio(model: Module) -> float:
             uncompressed_bits = get_torch_bit_depth(parameter)
             compressed_bits = uncompressed_bits
             if (
-                is_submodule_quantized(submodule)
+                is_model_quantized(submodule)
                 and submodule.quantization_scheme.weights
             ):
                 compressed_bits = submodule.quantization_scheme.weights.num_bits