From 1ba1ef0ec1d6a28efcaa7842f7bf5ac142ac49c0 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Fri, 23 Aug 2024 01:39:04 +0300 Subject: [PATCH] Update deepspeed_utils.py --- library/deepspeed_utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/library/deepspeed_utils.py b/library/deepspeed_utils.py index 99a7b2b..39583c7 100644 --- a/library/deepspeed_utils.py +++ b/library/deepspeed_utils.py @@ -91,9 +91,9 @@ def prepare_deepspeed_plugin(args: argparse.Namespace): zero3_save_16bit_model=args.zero3_save_16bit_model, ) deepspeed_plugin.deepspeed_config["train_micro_batch_size_per_gpu"] = args.train_batch_size - deepspeed_plugin.deepspeed_config["train_batch_size"] = ( - args.train_batch_size * args.gradient_accumulation_steps * int(os.environ["WORLD_SIZE"]) - ) + deepspeed_plugin.deepspeed_config["train_batch_size"] = 1#( + # args.train_batch_size * args.gradient_accumulation_steps * int(os.environ["WORLD_SIZE"]) + #) deepspeed_plugin.set_mixed_precision(args.mixed_precision) if args.mixed_precision.lower() == "fp16": deepspeed_plugin.deepspeed_config["fp16"]["initial_scale_power"] = 0 # preventing overflow.