use low_cpu_mem_usage with ds zero 1 or 2

use low_cpu_mem_usage when using deepspeed
2024-01-16 19:33:44 -05:00 · 2024-01-16 07:44:35 -05:00
1 changed files with 3 additions and 0 deletions
--- a/src/axolotl/utils/models.py
+++ b/src/axolotl/utils/models.py
@@ -357,6 +357,9 @@ def load_model(
    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
    if cfg.model_revision:
        model_kwargs["revision"] = cfg.model_revision
Author	SHA1	Message	Date
Wing Lian	1b33588f09	use low_cpu_mem_usage with ds zero 1 or 2	2024-01-16 19:33:44 -05:00
Wing Lian	1b59a3e698	use low_cpu_mem_usage when using deepspeed	2024-01-16 07:44:35 -05:00