Repository: huggingface/diffusers
File: src/diffusers/quantizers/torchao/torchao_quantizer.py (Python)

Code before:
            # about AffineQuantizedTensor
            module._parameters[tensor_name] = torch.nn.Parameter(param_value.to(device=target_device))
            if isinstance(module, nn.Linear):
                module.extra_repr = types.MethodType(_linear_extra_repr, module)
        else:
            # As we perform quantization here, the repr of linear layers is that of AQT, so we don't have to do it ourselves
            module._parameters[tensor_name] = torch.nn.Parameter(param_value).to(device=target_device)
            quantize_(module, self.quantization_config.get_apply_tensor_subclass())

    def get_cuda_warm_up_factor(self):
        """
        This factor is used in caching_allocator_warmup to determine how many bytes to pre-allocate for CUDA warmup.
        - A factor of 2 means we pre-allocate the full memory footprint of the model.
        - A factor of 4 means we pre-allocate half of that, and so on

        However, when using TorchAO, calculating memory usage with param.numel() * param.element_size() doesn't give
        the correct size for quantized weights (like int4 or int8) That's because TorchAO internally represents
        quantized tensors using subtensors and metadata, and the reported element_size() still corresponds to the
        torch_dtype not the actual bit-width of the quantized data.

        To correct for this:
        - Use a division factor of 8 for int4 weights
        - Use a division factor of 4 for int8 weights
        """
        # Original mapping for non-AOBaseConfig types
        map_to_target_dtype = {"int4_*": 8, "int8_*": 4, "float8*": 4}
        quant_type = self.quantization_config.quant_type
        for pattern, target_dtype in map_to_target_dtype.items():
            if fnmatch(quant_type, pattern):
                return target_dtype
        raise ValueError(f"Unsupported quant_type: {quant_type!r}")

    def _process_model_before_weight_loading(
        self,
        model: "ModelMixin",
        device_map,
        keep_in_fp32_modules: List[str] = [],
        **kwargs,
    ):
        self.modules_to_not_convert = self.quantization_config.modules_to_not_convert

        if not isinstance(self.modules_to_not_convert, list):
            self.modules_to_not_convert = [self.modules_to_not_convert]

        self.modules_to_not_convert.extend(keep_in_fp32_modules)

        # Extend `self.modules_to_not_convert` to keys that are supposed to be offloaded to `cpu` or `disk`
        if isinstance(device_map, dict) and len(device_map.keys()) > 1:
            keys_on_cpu = [key for key, value in device_map.items() if value in ["disk", "cpu"]]
            self.modules_to_not_convert.extend(keys_on_cpu)


Review comment (suggestion):
Need to handle more of these exhaustively: 

https://github.com/huggingface/diffusers/tree/1082c46afa4a15c49833d67c7f1c0f3cfd7b0570/src/diffusers/quantizers/quantization_config.py#L633

Code after:
            # about AffineQuantizedTensor
            module._parameters[tensor_name] = torch.nn.Parameter(param_value.to(device=target_device))
            if isinstance(module, nn.Linear):
                module.extra_repr = types.MethodType(_linear_extra_repr, module)
        else:
            # As we perform quantization here, the repr of linear layers is that of AQT, so we don't have to do it ourselves
            module._parameters[tensor_name] = torch.nn.Parameter(param_value).to(device=target_device)
            quantize_(module, self.quantization_config.get_apply_tensor_subclass())

    def get_cuda_warm_up_factor(self):
        """
        This factor is used in caching_allocator_warmup to determine how many bytes to pre-allocate for CUDA warmup.
        - A factor of 2 means we pre-allocate the full memory footprint of the model.
        - A factor of 4 means we pre-allocate half of that, and so on

        However, when using TorchAO, calculating memory usage with param.numel() * param.element_size() doesn't give
        the correct size for quantized weights (like int4 or int8) That's because TorchAO internally represents
        quantized tensors using subtensors and metadata, and the reported element_size() still corresponds to the
        torch_dtype not the actual bit-width of the quantized data.

        To correct for this:
        - Use a division factor of 8 for int4 weights
        - Use a division factor of 4 for int8 weights
        """
        # Original mapping for non-AOBaseConfig types
        # For the uint types, this is a best guess. Once these types become more used
        # we can look into their nuances.
        map_to_target_dtype = {"int4_*": 8, "int8_*": 4, "uint*": 8, "float8*": 4}
        quant_type = self.quantization_config.quant_type
        for pattern, target_dtype in map_to_target_dtype.items():
            if fnmatch(quant_type, pattern):
                return target_dtype
        raise ValueError(f"Unsupported quant_type: {quant_type!r}")

    def _process_model_before_weight_loading(
        self,
        model: "ModelMixin",
        device_map,
        keep_in_fp32_modules: List[str] = [],
        **kwargs,
    ):
        self.modules_to_not_convert = self.quantization_config.modules_to_not_convert

        if not isinstance(self.modules_to_not_convert, list):
            self.modules_to_not_convert = [self.modules_to_not_convert]

        self.modules_to_not_convert.extend(keep_in_fp32_modules)

        # Extend `self.modules_to_not_convert` to keys that are supposed to be offloaded to `cpu` or `disk`
        if isinstance(device_map, dict) and len(device_map.keys()) > 1:
            keys_on_cpu = [key for key, value in device_map.items() if value in ["disk", "cpu"]]