fix: Always return a list of quantization bits values from `get_quant…

…ization` (#3926)

merging– failing tests are unrelated

ludwig/utils/config_utils.py

@@ -144,7 +144,7 @@ def config_uses_llm(config: Union[Dict[str, Any], ModelConfig]) -> bool:
     return uses_llm
 
 
-def get_quantization(config: Union[Dict[str, Any], ModelConfig]) -> Union[int, List[int], None]:
+def get_quantization(config: Union[Dict[str, Any], ModelConfig]) -> List[Union[int, None]]:
     """Get the quantization specified in a config at any level.
 
     Args:
@@ -157,7 +157,7 @@ def get_quantization(config: Union[Dict[str, Any], ModelConfig]) -> Union[int, L
     """
     if isinstance(config, ModelConfig):
         if config.model_type == MODEL_LLM:
-            return config.quantization.bits if config.quantization else None
+            return [config.quantization.bits] if config.quantization else [None]
         else:
             quantization_bits = []
             for feature in config.input_features:
@@ -169,7 +169,7 @@ def get_quantization(config: Union[Dict[str, Any], ModelConfig]) -> Union[int, L
             return quantization_bits
     elif isinstance(config, dict) and config:
         if config.get(MODEL_TYPE, MODEL_ECD) == MODEL_LLM:
-            return config.get("quantization", {}).get("bits")
+            return [config.get("quantization", {}).get("bits")]
         elif INPUT_FEATURES in config:
             quantization_bits = []
             for feature in config.get(INPUT_FEATURES, []):

tests/ludwig/utils/test_config_utils.py

@@ -246,9 +246,9 @@ def ecd_config_dict_llm_encoder_8bit(
     "config,expectation",
     [
         # LLM configurations
-        ("llm_config_dict", None),
-        ("llm_config_dict_4bit", 4),
-        ("llm_config_dict_8bit", 8),
+        ("llm_config_dict", [None]),
+        ("llm_config_dict_4bit", [4]),
+        ("llm_config_dict_8bit", [8]),
         # LLM encoder configurations with one feature
         ("ecd_config_dict_llm_encoder", [None]),
         ("ecd_config_dict_llm_encoder_4bit", [4]),