fix: dont change quant storage dtype in case of fsdp (#1837)

* fix: dont change quant storage dtype in case of fsdp * fix black --------- Co-authored-by: Gal Cohen <[email protected]>
axolotl-ai-cloud · Aug 20, 2024 · 5aac4bc · 5aac4bc
1 parent e299312
commit 5aac4bc
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py
@@ -544,7 +544,9 @@ def load_model(
             "bnb_4bit_quant_type": "nf4",
             "bnb_4bit_quant_storage": torch.bfloat16,
         }
-        if cfg.model_config_type in ["jamba", "qwen2_moe"] and not cfg.deepspeed:
+        if cfg.model_config_type in ["jamba", "qwen2_moe"] and not (
+            cfg.deepspeed or cfg.fsdp
+        ):
             # for some reason, this causes the loss to be off by an order of magnitude
             # but deepspeed needs this still in bfloat16
             bnb_config["bnb_4bit_quant_storage"] = torch.float32