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

amend model loading for hqq + fix hqq version

Browse Source
This commit is contained in:
Sunny Liu
2025-04-21 15:53:29 -04:00
parent c8fb5baad6
commit f0a189131b
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1044,8 +1044,11 @@ class ModelLoader:
config=self.model_config, config=self.model_config,
) )
else: else:
if self.cfg.hqq: if self.cfg.hqq and torch.cuda.device_count() < 2:
# if using hqq, we need to set device_map to gpu otherwise the loading get stuck # for some reason on single gpu, we need to set device_map to auto/cuda
# otherwise you run into tensors on two devices error during training
# Doesn't affect multi-gpu tho
self.model_kwargs["device_map"] = "auto" self.model_kwargs["device_map"] = "auto"
self.model = self.auto_model_loader.from_pretrained( self.model = self.auto_model_loader.from_pretrained(
self.base_model, self.base_model,