tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

fix for fsdp for models that aren't qwen2 or jamba

Browse Source
This commit is contained in:
Wing Lian
2024-04-05 17:02:54 -07:00
parent bf4cd67252
commit 744f7082f5
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/axolotl/utils/models.py
View File

@@ -459,7 +459,7 @@ def load_model(
"bnb_4bit_quant_type": "nf4", "bnb_4bit_quant_type": "nf4",
"bnb_4bit_quant_storage": torch.bfloat16, "bnb_4bit_quant_storage": torch.bfloat16,
} }
if not cfg.deepspeed: if not cfg.deepspeed and cfg.model_config_type in ("jamba", "qwen2_moe"):
# for some reason, this causes the loss to be off by an order of magnitude # for some reason, this causes the loss to be off by an order of magnitude
# but deepspeed needs this still in bfloat16 # but deepspeed needs this still in bfloat16
bnb_config["bnb_4bit_quant_storage"] = torch.float32 bnb_config["bnb_4bit_quant_storage"] = torch.float32