From 2f670a5988a6083aa649b21ce67fb913ff7f450d Mon Sep 17 00:00:00 2001 From: mhenrhcsen Date: Wed, 16 Jul 2025 21:16:47 +0200 Subject: [PATCH] Fix: Update model loading logic to conditionally upcast based on lm_head presence for btlm models --- src/axolotl/loaders/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/axolotl/loaders/model.py b/src/axolotl/loaders/model.py index 7fce928f0..3d11601ba 100644 --- a/src/axolotl/loaders/model.py +++ b/src/axolotl/loaders/model.py @@ -798,7 +798,7 @@ class ModelLoader: if before_kbit_train_or_finetune: if name.endswith(".gate"): 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 continue if any(m in name for m in embedding_modules) and hasattr(module, "weight"):