Merge branch 'main' of https://github.com/microsoft/Megatron-DeepSpeed …

…into hzheng-data-fix
argonne-lcf · Oct 18, 2024 · ab3a8ec · ab3a8ec
2 parents 68da2db + deb95cd
commit ab3a8ec
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,