.

Author: jedmccaleb

megatron/training.py

@@ -613,7 +613,7 @@ def setup_model_and_optimizer(neox_args, use_cache=False, iteration=None):
 
         params=neox_args.deepspeed_config
         params["gradient_accumulation_steps"]=1 # JED figure out what is broken
-        params["train_batch_size"]=4 # JED figure out what is broken
+        #params["train_batch_size"]=4 # JED figure out what is broken
 
 
         model, optimizer, _, lr_scheduler = deepspeed.initialize(