fsdp requires params be the same type too (#493)
Browse files
src/axolotl/utils/models.py
CHANGED
|
@@ -356,7 +356,7 @@ def load_model(
|
|
| 356 |
if hasattr(module, "weight"):
|
| 357 |
module.to(torch.float32)
|
| 358 |
|
| 359 |
-
needs_fa2_dtype = cfg.adapter
|
| 360 |
if not cfg.gptq and (
|
| 361 |
(cfg.adapter == "lora" and load_in_8bit)
|
| 362 |
or (cfg.adapter == "qlora" and cfg.load_in_4bit)
|
|
|
|
| 356 |
if hasattr(module, "weight"):
|
| 357 |
module.to(torch.float32)
|
| 358 |
|
| 359 |
+
needs_fa2_dtype = cfg.adapter or cfg.fsdp
|
| 360 |
if not cfg.gptq and (
|
| 361 |
(cfg.adapter == "lora" and load_in_8bit)
|
| 362 |
or (cfg.adapter == "qlora" and cfg.load_in_4bit)
|