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

use low_cpu_mem_usage when using deepspeed

Browse Source
This commit is contained in:
Wing Lian
2024-01-16 07:44:35 -05:00
parent ece0211996
commit 1b59a3e698
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -355,8 +355,9 @@ def load_model(
# else: # else:
# model_kwargs["device_map"] = "cuda:" + str(torch.cuda.current_device()) # model_kwargs["device_map"] = "cuda:" + str(torch.cuda.current_device())
if is_deepspeed_zero3_enabled(): if is_deepspeed_zero3_enabled() or cfg.deepspeed:
del model_kwargs["device_map"] del model_kwargs["device_map"]
model_kwargs["low_cpu_mem_usage"] = True
if cfg.model_revision: if cfg.model_revision:
model_kwargs["revision"] = cfg.model_revision model_kwargs["revision"] = cfg.model_revision