Update deepspeed_utils.py

2024-08-23 01:39:04 +03:00
parent 9c4bba839a
commit 1ba1ef0ec1
1 changed files with 3 additions and 3 deletions
@@ -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.