use low_cpu_mem_usage when using deepspeed

This commit is contained in:
Wing Lian
2024-01-16 07:44:35 -05:00
parent ece0211996
commit 1b59a3e698

View File

@@ -355,8 +355,9 @@ def load_model(
# else:
# 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"]
model_kwargs["low_cpu_mem_usage"] = True
if cfg.model_revision:
model_kwargs["revision"] = cfg.model_revision