From 1b33588f099b27dc0e10536a3b96eaee70d022b1 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Tue, 16 Jan 2024 19:33:44 -0500 Subject: [PATCH] use low_cpu_mem_usage with ds zero 1 or 2 --- src/axolotl/utils/models.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py index 394e335f4..1a5836122 100644 --- a/src/axolotl/utils/models.py +++ b/src/axolotl/utils/models.py @@ -355,7 +355,9 @@ def load_model( # else: # model_kwargs["device_map"] = "cuda:" + str(torch.cuda.current_device()) - if is_deepspeed_zero3_enabled() or cfg.deepspeed: + if is_deepspeed_zero3_enabled(): + del model_kwargs["device_map"] + elif cfg.deepspeed: del model_kwargs["device_map"] model_kwargs["low_cpu_mem_usage"] = True