From 64873c1b4317afad99a1d397454ba0c64c6cb0b1 Mon Sep 17 00:00:00 2001 From: BootsofLagrangian Date: Mon, 5 Feb 2024 17:11:50 +0900 Subject: [PATCH] fix offload_optimizer_device typo --- library/train_util.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/library/train_util.py b/library/train_util.py index 2d85c977..933a34c4 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -3134,7 +3134,7 @@ def add_training_arguments(parser: argparse.ArgumentParser, support_dreambooth: help="Possible options are 0,1,2,3." ) parser.add_argument( - "--offload_optimizer", + "--offload_optimizer_device", type=str, default=None, choices=[None, "cpu", "nvme"], help="Possible options are none|cpu|nvme. Only applicable with ZeRO Stages 2 and 3." @@ -3959,7 +3959,7 @@ def prepare_accelerator(args: argparse.Namespace): deepspeed_plugin = DeepSpeedPlugin( zero_stage=args.zero_stage, gradient_accumulation_steps=args.gradient_accumulation_steps, gradient_clipping=args.max_grad_norm, - offload_optimizer=args.offload_optimizer, offload_optimizer_nvme_path=args.offload_optimizer_nvme_path, + offload_optimizer_device=args.offload_optimizer_device, offload_optimizer_nvme_path=args.offload_optimizer_nvme_path, offload_param_device=args.offload_param_device, offload_param_nvme_path=args.offload_param_nvme_path, zero3_init_flag=args.zero3_init_flag, zero3_save_16bit_model=args.zero3_save_16bit_model, )