diff --git a/ludwig/schema/llms/quantization.py b/ludwig/schema/llms/quantization.py index 2812f59afa9..c42871bc0de 100644 --- 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 index 563319d432b..299e28f4e80 100644 --- 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