[ZeRO-3] Ensured passing neox deepspeed_config when using partitioned…

… init (EleutherAI#1191) * added ds zero.Init() to get_model * Clean up conditional with block * pre-commit * ensured deepspeed configs are passed to init --------- Co-authored-by: Quentin Anthony <[email protected]>
CERC-AAI · Apr 1, 2024 · 01657aa · 01657aa
1 parent 51a7de9
commit 01657aa
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/megatron/training.py b/megatron/training.py
@@ -428,7 +428,9 @@ def get_model(neox_args, use_cache=False):
     old_use_mup = neox_args.use_mup
     neox_args.use_mup = False
 
-    with deepspeed.zero.Init() if neox_args.zero_stage == 3 else nullcontext() as gs:
+    with deepspeed.zero.Init(
+        config_dict_or_path=neox_args.deespeed_config
+    ) if neox_args.zero_stage == 3 else nullcontext() as gs:
         model = GPT2ModelPipe(
             neox_args=neox_args,
             num_tokentypes=0,