Merge pull request #768 from NVIDIA/am/megatron-cmd

Do not enable recompute-activations by default

Author: amaslenn

src/cloudai/workloads/megatron_run/megatron_run.py

@@ -1,5 +1,5 @@
 # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -36,7 +36,7 @@ class MegatronRunCmdArgs(CmdArgs):
     num_attention_heads: Optional[int] = 32
     num_layers: Optional[int] = 32
     pipeline_model_parallel_size: Optional[int] = 1
-    recompute_activations: Optional[str] = ""
+    recompute_activations: Optional[str] = None
     seq_length: Optional[int] = 4096
     tensor_model_parallel_size: Optional[int] = 2
 

tests/ref_data/megatron-run.sbatch

@@ -14,4 +14,4 @@ srun --export=ALL --mpi=pmix -N1 --container-image=nvcr.io/nvidia/megatron:24.09
 
 srun --export=ALL --mpi=pmix -N1 --container-image=nvcr.io/nvidia/megatron:24.09 --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output,$PWD --ntasks=1 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err bash /cloudai_install/slurm-metadata.sh
 
-srun --export=ALL --mpi=pmix -N1 --container-image=nvcr.io/nvidia/megatron:24.09 --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output,$PWD bash -c "source __OUTPUT_DIR__/output/env_vars.sh; python __CLOUDAI_DIR__/run.py --global-batch-size 16 --hidden-size 4096 --max-position-embeddings 4096 --num-attention-heads 32 --num-layers 32 --pipeline-model-parallel-size 1 --recompute-activations  --seq-length 4096 --tensor-model-parallel-size 2 --save __CLOUDAI_DIR__ --load __CLOUDAI_DIR__ --tokenizer-model __CLOUDAI_DIR__/model.m"
+srun --export=ALL --mpi=pmix -N1 --container-image=nvcr.io/nvidia/megatron:24.09 --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output,$PWD bash -c "source __OUTPUT_DIR__/output/env_vars.sh; python __CLOUDAI_DIR__/run.py --global-batch-size 16 --hidden-size 4096 --max-position-embeddings 4096 --num-attention-heads 32 --num-layers 32 --pipeline-model-parallel-size 1 --seq-length 4096 --tensor-model-parallel-size 2 --save __CLOUDAI_DIR__ --load __CLOUDAI_DIR__ --tokenizer-model __CLOUDAI_DIR__/model.m"

tests/test_test_definitions.py

@@ -1,5 +1,5 @@
 # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
-# Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -283,10 +283,15 @@ def test_default(self, megatron_run: MegatronRunTestDefinition):
         assert "--num-attention-heads 32" in cmd
         assert "--num-layers 32" in cmd
         assert "--pipeline-model-parallel-size 1" in cmd
-        assert "--recompute-activations " in cmd
+        assert "--recompute-activations" not in cmd
         assert "--seq-length 4096" in cmd
         assert "--tensor-model-parallel-size 2" in cmd
 
+    def test_recompute_activations_set(self, megatron_run: MegatronRunTestDefinition):
+        megatron_run.cmd_args.recompute_activations = ""
+        cmd = " ".join([f"{k} {v}" for k, v in megatron_run.cmd_args_dict.items()])
+        assert "--recompute-activations " in cmd
+
     def test_nones_are_dropped(self, megatron_run: MegatronRunTestDefinition):
         to_be_none = {
             "hidden_size": None,