diff --git a/modules/modelLoader/mixin/HFModelLoaderMixin.py b/modules/modelLoader/mixin/HFModelLoaderMixin.py index f2f196257..50ad05901 100644 --- a/modules/modelLoader/mixin/HFModelLoaderMixin.py +++ b/modules/modelLoader/mixin/HFModelLoaderMixin.py @@ -157,7 +157,13 @@ def __load_sub_module( if torch.is_floating_point(old_value): old_type = type(old_value) if not is_quantized_parameter(module, tensor_name): - if dtype.is_quantized() or module_name in keep_in_fp32_modules: + if module_name in keep_in_fp32_modules: + value = value.to(dtype=train_dtype.torch_dtype()) + elif dtype.is_quantized() and type(module) is nn.Linear: + # a plain Linear that the quantization layer filter excluded + fallback_dtype = quantization.fallback_dtype if quantization is not None else DataType.BFLOAT_16 + value = value.to(dtype=fallback_dtype.torch_dtype()) + elif dtype.is_quantized(): value = value.to(dtype=train_dtype.torch_dtype()) else: value = value.to(dtype=dtype.torch_dtype()) @@ -283,7 +289,13 @@ def __convert_sub_module_to_dtype( if value is not None and torch.is_floating_point(value): old_type = type(value) if not is_quantized_parameter(module, tensor_name): - if dtype.is_quantized() or module_name in keep_in_fp32_modules: + if module_name in keep_in_fp32_modules: + value = value.to(dtype=train_dtype.torch_dtype()) + elif dtype.is_quantized() and type(module) is nn.Linear: + # a plain Linear that the quantization layer filter excluded + fallback_dtype = quantization.fallback_dtype if quantization is not None else DataType.BFLOAT_16 + value = value.to(dtype=fallback_dtype.torch_dtype()) + elif dtype.is_quantized(): value = value.to(dtype=train_dtype.torch_dtype()) else: value = value.to(dtype=dtype.torch_dtype()) diff --git a/modules/ui/BaseModelTabView.py b/modules/ui/BaseModelTabView.py index 3b48a46f0..9589f4e86 100644 --- a/modules/ui/BaseModelTabView.py +++ b/modules/ui/BaseModelTabView.py @@ -151,15 +151,6 @@ def __create_base_components( has_vae: bool = False, include_compressed: bool = False, ) -> int: - if has_unet: - # unet weight dtype - self.components.label(frame, row, 3, "UNet Data Type", - tooltip="The unet weight data type") - self.components.options_kv(frame, row, 4, self.__create_dtype_options(include_a8=True, include_compressed=include_compressed), - ui_state, "unet.weight_dtype") - - row += 1 - if has_prior: if allow_override_prior: # prior model @@ -178,35 +169,22 @@ def __create_base_components( row += 1 - if has_transformer: - if allow_override_transformer: - # transformer model - self.components.label(frame, row, 0, "Override Transformer / GGUF", - tooltip="Can be used to override the transformer in the base model. Safetensors and GGUF files are supported, local and on Huggingface. If a GGUF file is used, the DataType must also be set to GGUF") - self.components.path_entry( - frame, row, 1, ui_state, "transformer.model_name", - mode="file", path_modifier=path_util.json_path_modifier - ) - - # transformer weight dtype - self.components.label(frame, row, 3, "Transformer Data Type", - tooltip="The transformer weight data type") - self.components.options_kv(frame, row, 4, self.__create_dtype_options(include_gguf=True, include_a8=True, include_compressed=include_compressed), - ui_state, "transformer.weight_dtype") - - row += 1 - - if has_unconditional_transformer: - # unconditional transformer weight dtype - self.components.label(frame, row, 3, "Unconditional Transformer Data Type", - tooltip="The weight data type of the unconditional transformer, used for the negative branch of CFG during sampling") - self.components.options_kv(frame, row, 4, self.__create_dtype_options(include_a8=True, include_compressed=include_compressed), - ui_state, "unconditional_transformer.weight_dtype") + if has_transformer and allow_override_transformer: + # transformer model + self.components.label(frame, row, 0, "Override Transformer / GGUF", + tooltip="Can be used to override the transformer in the base model. Safetensors and GGUF files are supported, local and on Huggingface. If a GGUF file is used, the DataType must also be set to GGUF") + self.components.path_entry( + frame, row, 1, ui_state, "transformer.model_name", + mode="file", path_modifier=path_util.json_path_modifier + ) row += 1 presets = controller.get_presets() + # Quantization Layer Filter (col 0/1) is a tall widget (preset row, custom entry row, regex row). + # UNet/Transformer Data Type and the quantization Fallback Data Type share this row too (col 3/4), + # lined up so the data type row matches the preset row, and the fallback row matches the entry row. self.components.label(frame, row, 0, "Quantization") self.components.layer_filter_entry(frame, row, 1, ui_state, preset_var_name="quantization.layer_filter_preset", presets=presets, @@ -219,7 +197,38 @@ def __create_base_components( frame_color="transparent", ) - # SVDQuant - create vertical grids to match the size of layer_filter_entry + if has_unet or has_transformer: + dtype_label_frame = self.components.inline_frame(frame, row, 3) + dtype_entry_frame = self.components.inline_frame(frame, row, 4) + + if has_unet: + self.components.label(dtype_label_frame, 0, 0, "UNet Data Type", + tooltip="The unet weight data type") + self.components.options_kv(dtype_entry_frame, 0, 0, self.__create_dtype_options(include_a8=True, include_compressed=include_compressed), + ui_state, "unet.weight_dtype") + else: + self.components.label(dtype_label_frame, 0, 0, "Transformer Data Type", + tooltip="The transformer weight data type") + self.components.options_kv(dtype_entry_frame, 0, 0, self.__create_dtype_options(include_gguf=True, include_a8=True, include_compressed=include_compressed), + ui_state, "transformer.weight_dtype") + + self.components.label(dtype_label_frame, 1, 0, "Fallback Data Type", + tooltip="The weight data type used for layers excluded by the quantization layer filter. Can itself be a quantized type.") + self.components.options_kv(dtype_entry_frame, 1, 0, self.__create_dtype_options(include_a8=True, include_compressed=include_compressed), + ui_state, "quantization.fallback_dtype") + + row += 1 + + if has_unconditional_transformer: + # unconditional transformer weight dtype + self.components.label(frame, row, 3, "Unconditional Transformer Data Type", + tooltip="The weight data type of the unconditional transformer, used for the negative branch of CFG during sampling") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(include_a8=True, include_compressed=include_compressed), + ui_state, "unconditional_transformer.weight_dtype") + + row += 1 + + # SVDQuant svd_label_frame, svd_entry_frame = self._make_svd_frames(frame, row) self.components.label(svd_label_frame, 0, 0, "SVDQuant", tooltip="What datatype to use for SVDQuant weights decomposition.") diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index ba11de9d7..53ed08438 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -349,6 +349,7 @@ class QuantizationConfig(BaseConfig): layer_filter: str layer_filter_preset: str layer_filter_regex: bool + fallback_dtype: DataType svd_dtype: DataType svd_rank: int cache_dir: str @@ -361,6 +362,7 @@ def default_values(): data.append(("layer_filter", "", str, False)) data.append(("layer_filter_preset", "full", str, False)) data.append(("layer_filter_regex", False, bool, False)) + data.append(("fallback_dtype", DataType.BFLOAT_16, DataType, False)) data.append(("svd_dtype", DataType.NONE, DataType, False)) data.append(("svd_rank", 16, int, False)) data.append(("cache_dir", None, str, True)) diff --git a/modules/util/quantization_util.py b/modules/util/quantization_util.py index d7229687d..dbedb681a 100644 --- a/modules/util/quantization_util.py +++ b/modules/util/quantization_util.py @@ -102,6 +102,7 @@ def __replace_linear_layers( construct_fn, keep_in_fp32_modules: list[str] | None = None, filters: list[ModuleFilter] | None = None, + fallback_construct_fn = None, copy_parameters: bool = False, name_prefix: str = "", visited_modules: set[int] | None = None, @@ -121,10 +122,12 @@ def __replace_linear_layers( if isinstance(parent_module, (nn.ModuleList, nn.Sequential, nn.ModuleDict)): for key, module in (parent_module.items() if isinstance(parent_module, nn.ModuleDict) else enumerate(parent_module)): if isinstance(module, convert_type): - if filters is not None and len(filters) > 0 and not any(f.matches(name_prefix) for f in filters): + matches = filters is None or len(filters) == 0 or any(f.matches(name_prefix) for f in filters) + fn = construct_fn if matches else fallback_construct_fn + if fn is None: continue - quant_linear = __create_linear_layer(construct_fn, module, copy_parameters) + quant_linear = __create_linear_layer(fn, module, copy_parameters) parent_module[key] = quant_linear del module elif id(module) not in visited_modules: @@ -133,6 +136,7 @@ def __replace_linear_layers( construct_fn=construct_fn, keep_in_fp32_modules=keep_in_fp32_modules, filters=filters, + fallback_construct_fn=fallback_construct_fn, copy_parameters=copy_parameters, name_prefix=f"{name_prefix}.{key}", visited_modules=visited_modules, @@ -145,10 +149,12 @@ def __replace_linear_layers( module = getattr(parent_module, attr_name) if isinstance(module, convert_type): key_name = attr_name if name_prefix == "" else f"{name_prefix}.{attr_name}" - if filters is not None and len(filters) > 0 and not any(f.matches(key_name) for f in filters): + matches = filters is None or len(filters) == 0 or any(f.matches(key_name) for f in filters) + fn = construct_fn if matches else fallback_construct_fn + if fn is None: continue - quant_linear = __create_linear_layer(construct_fn, module, copy_parameters) + quant_linear = __create_linear_layer(fn, module, copy_parameters) setattr(parent_module, attr_name, quant_linear) del module elif isinstance(module, nn.Module) and id(module) not in visited_modules: @@ -157,46 +163,51 @@ def __replace_linear_layers( construct_fn=construct_fn, keep_in_fp32_modules=keep_in_fp32_modules, filters=filters, + fallback_construct_fn=fallback_construct_fn, copy_parameters=copy_parameters, name_prefix=attr_name if name_prefix == "" else f"{name_prefix}.{attr_name}", visited_modules=visited_modules, ) -def replace_linear_with_quantized_layers( - parent_module: nn.Module, - dtype: DataType, - keep_in_fp32_modules: list[str] | None = None, - quantization: QuantizationConfig | None = None, - copy_parameters: bool = False, -): +def __quantized_linear_class_and_kwargs(dtype: DataType): + # deferred imports: the quantized Linear modules pull in heavy backends, so they stay + # out of module scope and are only imported once a quantized dtype is actually requested from modules.module.quantized.LinearFp8 import LinearFp8 from modules.module.quantized.LinearGGUFA8 import LinearGGUFA8 - from modules.module.quantized.LinearSVD import make_svd_linear from modules.module.quantized.LinearW8A8 import LinearW8A8 - kwargs = {} if dtype.quantize_nf4(): - linear_class = LinearNf4 + return LinearNf4, {} elif dtype.quantize_int8(): - linear_class = bnb.nn.Linear8bitLt - kwargs = {'has_fp16_weights': False} + return bnb.nn.Linear8bitLt, {'has_fp16_weights': False} elif dtype.quantize_fp8(): - linear_class = LinearFp8 + return LinearFp8, {} elif dtype.quantize_intW8A8(): - linear_class = LinearW8A8 - kwargs = {'dtype': torch.int8} + return LinearW8A8, {'dtype': torch.int8} elif dtype.quantize_fpW8A8(): - linear_class=LinearW8A8 - kwargs = {'dtype': torch.float8_e4m3fn} + return LinearW8A8, {'dtype': torch.float8_e4m3fn} elif dtype == DataType.GGUF_A8_INT: - linear_class=LinearGGUFA8 - kwargs = {'dtype': torch.int8} + return LinearGGUFA8, {'dtype': torch.int8} elif dtype == DataType.GGUF_A8_FLOAT: - linear_class=LinearGGUFA8 - kwargs = {'dtype': torch.float8_e4m3fn} + return LinearGGUFA8, {'dtype': torch.float8_e4m3fn} else: + return None, {} + +def replace_linear_with_quantized_layers( + parent_module: nn.Module, + dtype: DataType, + keep_in_fp32_modules: list[str] | None = None, + quantization: QuantizationConfig | None = None, + copy_parameters: bool = False, +): + from modules.module.quantized.LinearGGUFA8 import LinearGGUFA8 + from modules.module.quantized.LinearSVD import make_svd_linear + + linear_class, kwargs = __quantized_linear_class_and_kwargs(dtype) + if linear_class is None: return + fallback_construct_fn = None if quantization is not None: if quantization.svd_dtype != DataType.NONE: if dtype.is_gguf(): @@ -210,6 +221,11 @@ def replace_linear_with_quantized_layers( ModuleFilter(pattern, use_regex=quantization.layer_filter_regex) for pattern in quantization.layer_filter.split(",") ] + + if not dtype.is_gguf() and quantization.fallback_dtype.is_quantized(): + fallback_linear_class, fallback_kwargs = __quantized_linear_class_and_kwargs(quantization.fallback_dtype) + if fallback_linear_class is not None: + fallback_construct_fn = partial(fallback_linear_class, **fallback_kwargs) else: quant_filters = None @@ -219,6 +235,7 @@ def replace_linear_with_quantized_layers( construct_fn=partial(linear_class, **kwargs), keep_in_fp32_modules=keep_in_fp32_modules, filters=quant_filters, + fallback_construct_fn=fallback_construct_fn, copy_parameters=copy_parameters, convert_type=convert_type, )