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

Support --lora-on-cpu flag for DPO model merging (#2766) [skip ci]

Browse Source 
* Support --lora-on-cpu flag for DPO model merging

* fix: use device=cpu in _convert_embedding_modules_dtype when lora_on_cpu is set
This commit is contained in:
kallewoof
2025-06-28 00:19:24 +09:00
committed by GitHub
parent 0a7a216b60
commit ec15a7a691
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -776,6 +776,9 @@ class ModelLoader:
dist_dtype: torch.dtype, dist_dtype: torch.dtype,
before_kbit_train_or_finetune: bool, before_kbit_train_or_finetune: bool,
): ):
dest = {"dtype": dist_dtype}
if self.cfg.lora_on_cpu:
dest["device"] = "cpu"
for name, module in self.model.named_modules(): for name, module in self.model.named_modules():
if "norm" in name: if "norm" in name:
module.to(dist_dtype) module.to(dist_dtype)
@@ -786,4 +789,4 @@ class ModelLoader:
# 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"):
module.to(dist_dtype) module.to(**dest)