test fix

Author: ved1beta

src/axolotl/loaders/model.py

@@ -556,8 +556,9 @@ def _set_quantization_config(self):
             self.cfg.adapter in ["qlora", "qalora"]
             and self.model_kwargs["load_in_4bit"]
         ):
+            quantization_config = getattr(self.model_config, "quantization_config", {})
             self.model_kwargs["quantization_config"] = BitsAndBytesConfig(
-                **self.model_config.quantization_config
+                **quantization_config
             )
         elif self.cfg.adapter == "qlora" and self.model_kwargs.get(
             "load_in_4bit", False

src/axolotl/utils/schemas/config.py

@@ -541,6 +541,12 @@ class AxolotlInputConfig(
             "description": "Whether to use flash-attention rms norm implementation - advanced use only"
         },
     )
+    flash_attn_fuse_qkv: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Whether to fuse QKV projection into a single operation"
+        },
+    )
     flash_attn_fuse_mlp: bool | None = Field(
         default=None,
         json_schema_extra={