diff --git a/tests/functional_tests/shell_test_utils/start_interactive_job.sh b/tests/functional_tests/shell_test_utils/start_interactive_job.sh index 0b30fc01283..6d1b6f0c9bd 100644 --- a/tests/functional_tests/shell_test_utils/start_interactive_job.sh +++ b/tests/functional_tests/shell_test_utils/start_interactive_job.sh @@ -26,6 +26,7 @@ DATASET_DIR="" TIME="1:00:00" RECIPES_DIR="tests/test_utils/recipes" CONTAINER_MOUNTS="" +NO_GPUS_PER_TASK="FALSE" # Declare associative array for tracking unique mounts declare -A seen_mounts @@ -53,6 +54,10 @@ while [[ $# -gt 0 ]]; do TIME="$2" shift 2 ;; + --no-gpus-per-task) + NO_GPUS_PER_TASK="TRUE" + shift 1 + ;; --help) print_usage exit 0 @@ -89,7 +94,7 @@ SRUN_CMD="srun \ --container-workdir=/opt/megatron-lm \ --container-mounts=$CONTAINER_MOUNTS \ --nodes=1 \ - --gpus-per-task=8 \ + $(if [ "$NO_GPUS_PER_TASK" = "FALSE" ]; then echo "--gpus-per-task=8"; fi) \ --time=$TIME \ --pty bash" diff --git a/tests/test_utils/python_scripts/generate_local_jobs.py b/tests/test_utils/python_scripts/generate_local_jobs.py index 4a7cf2d7c13..88256cd93af 100644 --- a/tests/test_utils/python_scripts/generate_local_jobs.py +++ b/tests/test_utils/python_scripts/generate_local_jobs.py @@ -1,3 +1,5 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + """Generate launch scripts for local execution. This script allows to generate pre-filled launch scripts that allow for local execution of Megatron-LM functional tests inside containerized enviroments (i.e. Slurm enroot or Docker). @@ -24,7 +26,9 @@ def load_script(config_path: str) -> str: @click.command() @click.option("--model", required=False, type=str, help="Filters all tests by matching model") -@click.option("--scope", required=False, type=str, help="Filters all tests by matching scope") +@click.option( + "--scope", required=False, type=str, default="mr", help="Filters all tests by matching scope" +) @click.option( "--test-case", required=False, type=str, help="Returns a single test-case with matching name." ) @@ -77,8 +81,6 @@ def main( container_tag="none", ) - print(workloads) - for workload in workloads: if workload.type == "build": continue