Added quantization parameter metadata (#3487)

Co-authored-by: ksbrar <[email protected]>
ludwig-ai · Jul 31, 2023 · a09669e · a09669e
1 parent f8708a3
commit a09669e
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 2 deletions.
diff --git a/ludwig/schema/llms/quantization.py b/ludwig/schema/llms/quantization.py
@@ -4,6 +4,8 @@
 
 from ludwig.api_annotations import DeveloperAPI
 from ludwig.schema import utils as schema_utils
+from ludwig.schema.metadata import LLM_METADATA
+from ludwig.schema.metadata.parameter_metadata import convert_metadata_to_json
 from ludwig.schema.utils import ludwig_dataclass
 
 warnings.filterwarnings(
@@ -20,6 +22,7 @@ class QuantizationConfig(schema_utils.BaseMarshmallowConfig):
         options=[4, 8],
         default=4,
         description="The quantization level to apply to weights on load.",
+        parameter_metadata=LLM_METADATA["quantization"]["bits"],
     )
 
     llm_int8_threshold: float = schema_utils.NonNegativeFloat(
@@ -88,13 +91,20 @@ def __init__(self):
     def _jsonschema_type_mapping(self):
         return {
             "oneOf": [
-                {"type": "null", "title": "disabled", "description": "Disable quantization."},
+                {
+                    "type": "null",
+                    "title": "disabled",
+                    "description": "Disable quantization.",
+                    "parameter_metadata": convert_metadata_to_json(LLM_METADATA["quantization"]["_oneOf"]["none"]),
+                },
                 {
                     **schema_utils.unload_jsonschema_from_marshmallow_class(QuantizationConfig),
                     "title": "enabled",
                     "description": "Set quantization options.",
+                    "parameter_metadata": convert_metadata_to_json(LLM_METADATA["quantization"]["_oneOf"]["object"]),
                 },
             ],
             "title": "quantization",
-            "description": "",
+            "description": "Set quantization options.",
+            "parameter_metadata": convert_metadata_to_json(LLM_METADATA["quantization"]["_meta"]),
         }
diff --git a/ludwig/schema/metadata/configs/llm.yaml b/ludwig/schema/metadata/configs/llm.yaml
@@ -119,3 +119,17 @@ adapter:
     adapter_layers:
       ui_display_name: Adapter Layers
       expected_impact: 3
+quantization:
+  _oneOf:
+    object:
+      ui_display_name: Quantization
+      expected_impact: 3
+    none:
+      ui_display_name: No Quantization
+      expected_impact: 3
+  _meta:
+    expected_impact: 3
+    ui_component_type: radio_string_combined
+  bits:
+    ui_display_name: Bits per parameter
+    expected_impact: 3