fix init issue for silently ignoring the deepspeed config

microsoft · Oct 15, 2024 · 008b413 · 008b413
1 parent 4448492
commit 008b413
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/megatron/initialize.py b/megatron/initialize.py
@@ -185,6 +185,7 @@ def setup_deepspeed_random_and_activation_checkpointing(args):
 
     deepspeed.checkpointing.configure(
         mpu,
+        deepspeed_config=args.deepspeed_config,
         partition_activations=args.partition_activations,
         contiguous_checkpointing=args.contigious_checkpointing,
         num_checkpoints=num_layers,