Fix act config exporting for mixed schemes

auto_round/export/export_to_autoround/export.py

@@ -18,6 +18,7 @@
 import json
 import os
 from concurrent.futures import ThreadPoolExecutor
+from dataclasses import fields
 from enum import Enum
 
 import threadpoolctl as tctl
@@ -26,9 +27,10 @@
 import transformers
 from tqdm import tqdm
 
-from auto_round.export.export_to_autoround.utils import REQUIRED_CONFIG_KEYS, check_neq_config
+from auto_round.export.export_to_autoround.utils import check_neq_config
 from auto_round.export.utils import save_model
 from auto_round.logger import logger
+from auto_round.schemes import QuantizationScheme
 from auto_round.utils import (
     SUPPORTED_FORMATS,
     SUPPORTED_LAYER_TYPES,
@@ -324,26 +326,20 @@ def save_quantized_as_autoround(output_dir, inplace=True, backend="auto_round:ex
         for i in range(len(block_name_to_quantize)):
             block_name_to_quantize[i] = os.path.commonprefix(block_name_to_quantize[i]).rstrip(".")
 
-    for layer_name in layer_config:
-        if (
-            not layer_config[layer_name]["in_blocks"] and layer_config[layer_name]["bits"] <= 8
-        ):  ##lm head ##TODO fix act and so on
-            extra_config[layer_name] = {}
-            extra_config[layer_name]["bits"] = layer_config[layer_name]["bits"]
-            extra_config[layer_name]["data_type"] = layer_config[layer_name]["data_type"]
-            extra_config[layer_name]["group_size"] = layer_config[layer_name]["group_size"]
-            extra_config[layer_name]["sym"] = layer_config[layer_name]["sym"]
-        elif layer_config[layer_name]["in_blocks"] or (
+    scheme_keys = [f.name for f in fields(QuantizationScheme)]
+    for layer_name, cfg in layer_config.items():
+        if not cfg["in_blocks"] and cfg["bits"] <= 8:  # lm head
+            extra_config[layer_name] = {key: cfg.get(key) for key in scheme_keys}
+        elif cfg["in_blocks"] or (
             block_name_to_quantize is not None and check_start_with_block_name(layer_name, block_name_to_quantize)
         ):
-            neq_keys = check_neq_config(
-                layer_config[layer_name], **{k: quantization_config[k] for k in REQUIRED_CONFIG_KEYS}
-            )
+            neq_keys = check_neq_config(cfg, **{k: quantization_config[k] for k in scheme_keys})
             if len(neq_keys) > 0:
                 extra_config[layer_name] = {}
-            for key in neq_keys:
-                if layer_config[layer_name][key] is not None:
-                    extra_config[layer_name][key] = layer_config[layer_name][key]
+                for key in scheme_keys:
+                    if cfg[key] is not None:
+                        extra_config[layer_name][key] = cfg[key]
+
     if len(extra_config) > 0:
         quantization_config["extra_config"] = extra_config
     names = list(layer_config.keys())

auto_round/export/export_to_autoround/export_to_fp8.py

@@ -16,16 +16,18 @@
 import json
 import os
 from concurrent.futures import ThreadPoolExecutor
+from dataclasses import fields
 
 import threadpoolctl as tctl
 import torch
 import transformers
 from tqdm import tqdm
 
 from auto_round.data_type.utils import reshape_pad_tensor_by_group_size, revert_tensor_by_pad
-from auto_round.export.export_to_autoround.utils import REQUIRED_CONFIG_KEYS, check_neq_config
+from auto_round.export.export_to_autoround.utils import check_neq_config
 from auto_round.export.utils import save_model
 from auto_round.logger import logger
+from auto_round.schemes import QuantizationScheme
 from auto_round.utils import (
     SUPPORTED_LAYER_TYPES,
     _get_packing_device,
@@ -169,26 +171,20 @@ def save_quantized_as_autoround(output_dir, inplace=True, backend="auto_round",
         for i in range(len(block_name_to_quantize)):
             block_name_to_quantize[i] = os.path.commonprefix(block_name_to_quantize[i]).rstrip(".")
 
-    for layer_name in layer_config:
-        if (
-            not layer_config[layer_name]["in_blocks"] and layer_config[layer_name]["bits"] <= 8
-        ):  ##lm head ##TODO fix act and so on
-            extra_config[layer_name] = {}
-            extra_config[layer_name]["bits"] = layer_config[layer_name]["bits"]
-            extra_config[layer_name]["data_type"] = layer_config[layer_name]["data_type"]
-            extra_config[layer_name]["group_size"] = layer_config[layer_name]["group_size"]
-            extra_config[layer_name]["sym"] = layer_config[layer_name]["sym"]
-        elif layer_config[layer_name]["in_blocks"] or (
+    scheme_keys = [f.name for f in fields(QuantizationScheme)]
+    for layer_name, cfg in layer_config.items():
+        if not cfg["in_blocks"] and cfg["bits"] <= 8:  # lm head
+            extra_config[layer_name] = {key: cfg.get(key) for key in scheme_keys}
+        elif cfg["in_blocks"] or (
             block_name_to_quantize is not None and check_start_with_block_name(layer_name, block_name_to_quantize)
         ):
-            neq_keys = check_neq_config(
-                layer_config[layer_name], **{k: quantization_config[k] for k in REQUIRED_CONFIG_KEYS}
-            )
+            neq_keys = check_neq_config(cfg, **{k: quantization_config[k] for k in scheme_keys})
             if len(neq_keys) > 0:
                 extra_config[layer_name] = {}
-            for key in neq_keys:
-                if layer_config[layer_name][key] is not None:
-                    extra_config[layer_name][key] = layer_config[layer_name][key]
+                for key in scheme_keys:
+                    if cfg[key] is not None:
+                        extra_config[layer_name][key] = cfg[key]
+
     if len(extra_config) > 0:
         quantization_config["extra_config"] = extra_config
     names = list(layer_config.keys())

auto_round/export/export_to_autoround/export_to_nvfp_mxfp.py

@@ -17,16 +17,18 @@
 import json
 import os
 from concurrent.futures import ThreadPoolExecutor
+from dataclasses import fields
 
 import threadpoolctl as tctl
 import torch
 import torch.nn as nn
 import transformers
 from tqdm import tqdm
 
-from auto_round.export.export_to_autoround.utils import REQUIRED_CONFIG_KEYS, check_neq_config
+from auto_round.export.export_to_autoround.utils import check_neq_config
 from auto_round.export.utils import save_model
 from auto_round.logger import logger
+from auto_round.schemes import QuantizationScheme
 from auto_round.utils import (
     SUPPORTED_LAYER_TYPES,
     _get_packing_device,
@@ -195,26 +197,20 @@ def save_quantized_as_fp(output_dir, inplace=True, **kwargs):
         for i in range(len(block_name_to_quantize)):
             block_name_to_quantize[i] = os.path.commonprefix(block_name_to_quantize[i]).rstrip(".")
 
-    for layer_name in layer_config:
-        if (
-            not layer_config[layer_name]["in_blocks"] and layer_config[layer_name]["bits"] <= 8
-        ):  ##lm head # TODO fix act and so on
-            extra_config[layer_name] = {}
-            extra_config[layer_name]["bits"] = layer_config[layer_name]["bits"]
-            extra_config[layer_name]["data_type"] = layer_config[layer_name]["data_type"]
-            extra_config[layer_name]["group_size"] = layer_config[layer_name]["group_size"]
-            extra_config[layer_name]["sym"] = layer_config[layer_name]["sym"]
-        elif layer_config[layer_name]["in_blocks"] or (
+    scheme_keys = [f.name for f in fields(QuantizationScheme)]
+    for layer_name, cfg in layer_config.items():
+        if not cfg["in_blocks"] and cfg["bits"] <= 8:  # lm head
+            extra_config[layer_name] = {key: cfg.get(key) for key in scheme_keys}
+        elif cfg["in_blocks"] or (
             block_name_to_quantize is not None and check_start_with_block_name(layer_name, block_name_to_quantize)
         ):
-            neq_keys = check_neq_config(
-                layer_config[layer_name], **{k: quantization_config[k] for k in REQUIRED_CONFIG_KEYS}
-            )
+            neq_keys = check_neq_config(cfg, **{k: quantization_config[k] for k in scheme_keys})
             if len(neq_keys) > 0:
                 extra_config[layer_name] = {}
-            for key in neq_keys:
-                if layer_config[layer_name][key] is not None:
-                    extra_config[layer_name][key] = layer_config[layer_name][key]
+                for key in scheme_keys:
+                    if cfg[key] is not None:
+                        extra_config[layer_name][key] = cfg[key]
+
     if len(extra_config) > 0:
         quantization_config["extra_config"] = extra_config
     names = list(layer_config.keys())

auto_round/export/export_to_autoround/utils.py

@@ -12,36 +12,30 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-REQUIRED_CONFIG_KEYS = (
-    "data_type",
-    "bits",
-    "group_size",
-    "sym",
-    "act_bits",
-    "act_data_type",
-    "act_group_size",
-    "act_sym",
-    "act_dynamic",
-)
+from dataclasses import fields
+from typing import List
 
+from auto_round.schemes import QuantizationScheme
 
-def check_neq_config(config: dict, **expected) -> dict[str, tuple]:
+
+def check_neq_config(config: dict, **expected) -> List[str]:
     """
     Compare a config dict against expected values.
     Ensures all required keys are present in both config and expected.
 
     Returns:
-        dict[str, tuple]: {key: (actual, expected)} for mismatched values.
+        List[str]: [keys] for mismatched values.
     """
+    scheme_keys = [f.name for f in fields(QuantizationScheme)]
     # 1. Check missing from expected
-    missing_expected = [k for k in REQUIRED_CONFIG_KEYS if k not in expected]
+    missing_expected = [k for k in scheme_keys if k not in expected]
     if missing_expected:
         raise ValueError(f"Missing expected values for keys: {missing_expected}")
 
     # 2. Check missing from layer config
-    missing_config = [k for k in REQUIRED_CONFIG_KEYS if k not in config]
+    missing_config = [k for k in scheme_keys if k not in config]
     if missing_config:
         raise ValueError(f"Missing config values for keys: {missing_config}")
 
     # 3. Collect mismatches
-    return {key: (config[key], expected[key]) for key in REQUIRED_CONFIG_KEYS if config[key] != expected[key]}
+    return [key for key in scheme_keys if config[key] != expected[key] and config[key] is not None]