fallback dtype for quantization layer filter

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())

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.")

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))

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,
     )