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

Fix: Update model loading logic to conditionally upcast based on lm_head presence for btlm models

Browse Source
This commit is contained in:
mhenrhcsen
2025-07-16 21:16:47 +02:00
parent 84ad69afad
commit 2f670a5988
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/axolotl/loaders/model.py
View File

@@ -798,7 +798,7 @@ class ModelLoader:
if before_kbit_train_or_finetune: if before_kbit_train_or_finetune:
if name.endswith(".gate"): if name.endswith(".gate"):
module.to(dist_dtype) module.to(dist_dtype)
if self.model_config.model_type == "btlm": if self.model_config.model_type == "btlm" and "lm_head" in name:
# don't upcast lm_head for btlm # don't upcast lm_head for btlm
continue continue
if any(m in name for m in embedding_modules) and hasattr(module, "weight"): if any(m in name for m in embedding_modules) and hasattr(module, "weight"):