[core] Fix torchao (#42289)

* inital commit

* up

* update unexpected later on

* fix

* update

* simplify our lives

* isolate a bit more

* fixup

* small nits

* style

* nit

* fix common cases

* fix post merge

* bnb needs missing keys

* small fix

* bettrer documentation

* no veradict + base class

* rake review comments

* take all comments

* fix super init

* update doc to be more real

* up

* fix some tests

* weight convertor

* up

* mostly correct

* oups

* skip non linears

* only some tests to go

* need quantization

* fix tests

* rm comment

* revert

* revert 2

* style

* up

* up

* remove unsafe loading path

* fix

* fix

* fix

* up

* rm Dtensor import

* rm replicate import

* fix imports

* up

* minor modifications

* add quantizaton_operation

* delattr

* fix

* fix get_param_buffer

* better to just set module initialized

* rm tie_weights

* guard imports

* up

* rm offloading for now

* add license

* don't guard torch

* comment

* fix

* rm torch.grad

* revert

* fix

* add guard

* add second guard

---------

Co-authored-by: Arthur <[email protected]>
Co-authored-by: Marc Sun <[email protected]>

Author: 3 people

src/transformers/core_model_loading.py

@@ -115,6 +115,8 @@ def convert(
         source_keys: list[str],
         target_keys: list[str],
         full_layer_name: str,
+        model,
+        missing_keys,
         config,
         **kwargs,
     ) -> dict[str, list[torch.Tensor]]:
@@ -138,6 +140,8 @@ def convert(
         source_keys: list[str],
         target_keys: list[str],
         full_layer_name: str,
+        model,
+        missing_keys,
         config,
     ) -> dict[str, list[torch.Tensor]]:
         tensors = next(iter(value.values()))
@@ -163,6 +167,8 @@ def convert(
         source_keys: list[str],
         target_keys: list[str],
         full_layer_name: str,
+        model,
+        missing_keys,
         config,
     ) -> dict[str, torch.Tensor]:
         if len(target_keys) != 1:
@@ -191,6 +197,8 @@ def convert(
         source_keys: list[str],
         target_keys: list[str],
         full_layer_name: str,
+        model,
+        missing_keys,
         config,
     ) -> dict[str, torch.Tensor]:
         merged: dict[str, torch.Tensor] = {}
@@ -220,6 +228,8 @@ def convert(
         source_keys: list[str],
         target_keys: list[str],
         full_layer_name: str,
+        model,
+        missing_keys,
         config,
     ) -> dict[str, list[torch.Tensor]]:
         if len(value) != len(self.sizes):
@@ -258,6 +268,8 @@ def convert(
         source_keys: list[str],
         target_keys: list[str],
         full_layer_name: str,
+        model,
+        missing_keys,
         config,
     ) -> dict[str, list[torch.Tensor]]:
         self.config = config
@@ -298,21 +310,28 @@ def add_tensor(self, target_key: str, source_key: str, source_pattern: str, futu
 class WeightRenaming(WeightTransform):
     # Special case of WeightTransform that only renames keys without any conversion.
 
-    def convert(self, layer_name: str, config=None, quantizer=None, missing_keys: Optional[MutableSet[str]] = None):
+    def convert(
+        self,
+        layer_name: str,
+        model=None,
+        config=None,
+        hf_quantizer=None,
+        missing_keys: Optional[MutableSet[str]] = None,
+    ):
         misc = {}
         for pattern, futures in self.collected_tensors.items():
             self.collected_tensors[pattern] = [future.result() for future in futures]
 
         collected_tensors = self.collected_tensors
-        if quantizer is not None and self.quantization_operation is not None:
+        if hf_quantizer is not None and self.quantization_operation is not None:
             with log_to_misc(layer_name, misc, (self.collected_tensors, layer_name), self.quantization_operation):
                 collected_tensors = self.quantization_operation.convert(
                     self.collected_tensors,
                     source_keys=self.source_keys,
                     target_keys=self.target_keys,
                     full_layer_name=layer_name,
+                    model=model,
                     config=config,
-                    quant_config=quantizer.quantization_config,
                     missing_keys=missing_keys,
                 )
 
@@ -332,7 +351,14 @@ def __post_init__(self):
         if not self.operations:
             raise ValueError("WeightConverter requires at least one operation.")
 
-    def convert(self, layer_name: str, config=None, quantizer=None, missing_keys: Optional[MutableSet[str]] = None):
+    def convert(
+        self,
+        layer_name: str,
+        model=None,
+        config=None,
+        hf_quantizer=None,
+        missing_keys: Optional[MutableSet[str]] = None,
+    ):
         misc = {}
         for pattern, futures in self.collected_tensors.items():
             self.collected_tensors[pattern] = [future.result() for future in futures]
@@ -345,17 +371,19 @@ def convert(self, layer_name: str, config=None, quantizer=None, missing_keys: Op
                     source_keys=self.source_keys,
                     target_keys=self.target_keys,
                     full_layer_name=layer_name,
+                    model=model,
                     config=config,
+                    missing_keys=missing_keys,
                 )
-        if quantizer is not None and self.quantization_operation is not None:
+        if hf_quantizer is not None and self.quantization_operation is not None:
             with log_to_misc(layer_name, misc, (collected_tensors, layer_name), self.quantization_operation):
                 collected_tensors = self.quantization_operation.convert(
                     collected_tensors,
                     source_keys=self.source_keys,
                     target_keys=self.target_keys,
                     full_layer_name=layer_name,
                     config=config,
-                    quant_config=quantizer.quantization_config,
+                    model=model,
                     missing_keys=missing_keys,
                 )
         return collected_tensors, misc
@@ -626,7 +654,6 @@ def convert_and_load_state_dict_in_model(
     ```
 
     """
-
     prefix = model.base_model_prefix
     tp_plan = tp_plan or {}
     device_map = device_map or {"": "cpu"}
@@ -750,7 +777,11 @@ def convert_and_load_state_dict_in_model(
             pbar.refresh()
             try:
                 realized_value, misc = mapping.convert(
-                    first_param_name, config=model.config, quantizer=hf_quantizer, missing_keys=missing_keys
+                    first_param_name,
+                    model=model,
+                    config=model.config,
+                    hf_quantizer=hf_quantizer,
+                    missing_keys=missing_keys,
                 )
                 for target_name, param in realized_value.items():
                     param = param[0] if isinstance(param, list) else param

src/transformers/integrations/accelerate.py

@@ -241,7 +241,7 @@ def all_tensors():
         if name in tied_keys:
             continue
         if hf_quantizer is not None:
-            dtype_size = hf_quantizer.param_element_size(model, name)
+            dtype_size = hf_quantizer.param_element_size(model, name, param)
         else:
             dtype_size = param.element_size()
         size = param.numel() * dtype_size

src/transformers/integrations/bitsandbytes.py

@@ -36,7 +36,11 @@ def __init__(self, hf_quantizer):
         self.hf_quantizer = hf_quantizer
 
     def convert(
-        self, input_dict: torch.Tensor, model: Optional[torch.nn.Module] = None, missing_keys=None, **kwargs
+        self,
+        input_dict: dict[str, list[torch.Tensor]],
+        model: Optional[torch.nn.Module] = None,
+        missing_keys=None,
+        **kwargs,
     ) -> dict[str, torch.Tensor]:
         """
         we need to store some parameters to create the quantized weight. For example, bnb requires 6 values that are stored in the checkpoint to recover the quantized weight. So we store them in a dict that it stored in hf_quantizer for now as we can't save it in the op since we create an op per tensor.
@@ -59,6 +63,7 @@ def convert(
             # remove missing keys that were create when initializing Params4bit
             for key in new_value.quant_state.as_dict(packed=True).keys():
                 missing_keys.discard(f"{full_name}.{key}")
+            module._is_hf_initialized = True
             return {target_key: new_value}
         else:
             module_name = target_key.rsplit(".", 1)[0]
@@ -77,6 +82,7 @@ def convert(
                     device=value.device,
                     module=module,
                 )
+                module._is_hf_initialized = True
                 del self.hf_quantizer.param_quant_stats[module_name]
                 return {target_key: new_value}
             return {}
@@ -87,7 +93,11 @@ def __init__(self, hf_quantizer):
         self.hf_quantizer = hf_quantizer
 
     def convert(
-        self, input_dict: torch.Tensor, model: Optional[torch.nn.Module] = None, missing_keys=None, **kwargs
+        self,
+        input_dict: dict[str, list[torch.Tensor]],
+        model: Optional[torch.nn.Module] = None,
+        missing_keys=None,
+        **kwargs,
     ) -> dict[str, torch.Tensor]:
         target_key, value = tuple(input_dict.items())[0]
         value = value[0] if isinstance(value, list) else value

src/transformers/integrations/tensor_parallel.py

@@ -20,9 +20,13 @@
 from functools import partial, reduce
 from typing import Optional
 
-import torch
-import torch.distributed as dist
-from torch import nn
+from ..utils.import_utils import is_torch_available
+
+
+if is_torch_available():
+    import torch
+    import torch.distributed as dist
+    from torch import nn
 
 from ..distributed import DistributedConfig
 from ..utils import is_torch_greater_or_equal, logging
@@ -31,12 +35,12 @@
 
 logger = logging.get_logger(__name__)
 
-# Cache this result has it's a C FFI call which can be pretty time-consuming
-_torch_distributed_available = torch.distributed.is_available()
-
+if is_torch_available():
+    # Cache this result has it's a C FFI call which can be pretty time-consuming
+    _torch_distributed_available = torch.distributed.is_available()
 
-if is_torch_greater_or_equal("2.5") and _torch_distributed_available:
-    from torch.distributed.tensor import DTensor, Placement, Replicate, Shard
+    if is_torch_greater_or_equal("2.5") and _torch_distributed_available:
+        from torch.distributed.tensor import DTensor, Placement, Replicate, Shard
 
 
 def initialize_tensor_parallelism(
@@ -169,19 +173,20 @@ def _get_parameter_tp_plan(parameter_name: str, tp_plan: dict[str, str], is_weig
     return None
 
 
-str_to_dtype = {
-    "BOOL": torch.bool,
-    "U8": torch.uint8,
-    "I8": torch.int8,
-    "I16": torch.int16,
-    "F16": torch.float16,
-    "BF16": torch.bfloat16,
-    "I32": torch.int32,
-    "F32": torch.float32,
-    "F64": torch.float64,
-    "I64": torch.int64,
-    "F8_E4M3": torch.float8_e4m3fn,
-}
+if is_torch_available():
+    str_to_dtype = {
+        "BOOL": torch.bool,
+        "U8": torch.uint8,
+        "I8": torch.int8,
+        "I16": torch.int16,
+        "F16": torch.float16,
+        "BF16": torch.bfloat16,
+        "I32": torch.int32,
+        "F32": torch.float32,
+        "F64": torch.float64,
+        "I64": torch.int64,
+        "F8_E4M3": torch.float8_e4m3fn,
+    }
 
 
 def get_packed_weights(param, empty_param, device_mesh, rank, dim):