diff --git a/tests/functional_tests/shell_test_utils/_run_training.sh b/tests/functional_tests/shell_test_utils/_run_training.sh index 8f848a24add..654faac31d4 100644 --- a/tests/functional_tests/shell_test_utils/_run_training.sh +++ b/tests/functional_tests/shell_test_utils/_run_training.sh @@ -159,10 +159,13 @@ MASTER_PORT=${MASTER_PORT:-6000} NUM_NODES=${NUM_NODES:-${SLURM_NNODES:-1}} GPUS_PER_NODE=${GPUS_PER_NODE:-8} NODE_RANK=${SLURM_NODEID:-${SLURM_NODEID:-0}} -LAST_RANK=$((GPUS_PER_NODE - 1)) +LAST_RANK=$((GPUS_PER_NODE - 1)) export LOG_DIR=$OUTPUT_PATH/logs/$REPEAT mkdir -p $LOG_DIR +# Read launcher type from model config (default: torchrun) +LAUNCHER=$(/usr/local/bin/yq '.LAUNCHER // "torchrun"' "$TRAINING_PARAMS_PATH") + DISTRIBUTED_ARGS=( --nproc_per_node $GPUS_PER_NODE --nnodes $NUM_NODES @@ -176,11 +179,21 @@ DISTRIBUTED_ARGS=( # Start training if [[ "$IS_NEMO_TEST" == "true" ]]; then - uv run --no-sync python -m torch.distributed.run ${DISTRIBUTED_ARGS[@]} \ - --no-python /opt/venv/bin/$TRAINING_SCRIPT_PATH "${PARAMS[@]}" && EXIT_CODE=0 || EXIT_CODE=$? + if [[ "$LAUNCHER" == "ft_launcher" ]]; then + ft_launcher ${DISTRIBUTED_ARGS[@]} \ + --no-python /opt/venv/bin/$TRAINING_SCRIPT_PATH "${PARAMS[@]}" && EXIT_CODE=0 || EXIT_CODE=$? + else + uv run --no-sync python -m torch.distributed.run ${DISTRIBUTED_ARGS[@]} \ + --no-python /opt/venv/bin/$TRAINING_SCRIPT_PATH "${PARAMS[@]}" && EXIT_CODE=0 || EXIT_CODE=$? + fi else - uv run --no-sync python -m torch.distributed.run ${DISTRIBUTED_ARGS[@]} \ - $TRAINING_SCRIPT_PATH "${PARAMS[@]}" && EXIT_CODE=0 || EXIT_CODE=$? + if [[ "$LAUNCHER" == "ft_launcher" ]]; then + ft_launcher ${DISTRIBUTED_ARGS[@]} \ + $TRAINING_SCRIPT_PATH "${PARAMS[@]}" && EXIT_CODE=0 || EXIT_CODE=$? + else + uv run --no-sync python -m torch.distributed.run ${DISTRIBUTED_ARGS[@]} \ + $TRAINING_SCRIPT_PATH "${PARAMS[@]}" && EXIT_CODE=0 || EXIT_CODE=$? + fi fi # Run after script diff --git a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp1_pp4_vp1_resume_torch_dist_dist_optimizer_overlap_grad_reduce_param_gather/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp1_pp4_vp1_resume_torch_dist_dist_optimizer_overlap_grad_reduce_param_gather/model_config.yaml index bc18bf3e2d8..4f2d19ce559 100644 --- a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp1_pp4_vp1_resume_torch_dist_dist_optimizer_overlap_grad_reduce_param_gather/model_config.yaml +++ b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp1_pp4_vp1_resume_torch_dist_dist_optimizer_overlap_grad_reduce_param_gather/model_config.yaml @@ -61,3 +61,4 @@ MODEL_ARGS: --async-save: true --use-persistent-ckpt-worker: true TEST_TYPE: regular # Usually ckpt-resume, but as a WAR to #513 set to regular +LAUNCHER: ft_launcher diff --git a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_resume_torch_dist_multi_dist_optimizer_instances/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_resume_torch_dist_multi_dist_optimizer_instances/model_config.yaml index 42c361d5f8a..d226ebb89c5 100644 --- a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_resume_torch_dist_multi_dist_optimizer_instances/model_config.yaml +++ b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp1_resume_torch_dist_multi_dist_optimizer_instances/model_config.yaml @@ -57,3 +57,4 @@ MODEL_ARGS: --async-save: true --use-persistent-ckpt-worker: true TEST_TYPE: ckpt-resume +LAUNCHER: ft_launcher diff --git a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp4_pp1_resume_torch_dist_dist_optimizer_overlap_grad_reduce_param_gather/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp4_pp1_resume_torch_dist_dist_optimizer_overlap_grad_reduce_param_gather/model_config.yaml index 6cc70da70fc..12acfaf8402 100644 --- a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp4_pp1_resume_torch_dist_dist_optimizer_overlap_grad_reduce_param_gather/model_config.yaml +++ b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp4_pp1_resume_torch_dist_dist_optimizer_overlap_grad_reduce_param_gather/model_config.yaml @@ -57,3 +57,4 @@ MODEL_ARGS: --async-strategy: mcore --use-persistent-ckpt-worker: true TEST_TYPE: ckpt-resume +LAUNCHER: ft_launcher diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp1_pp2_resume_torch_dist_reshard_2x1x4_te_8experts2parallel_dist_optimizer/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp1_pp2_resume_torch_dist_reshard_2x1x4_te_8experts2parallel_dist_optimizer/model_config.yaml index 329639ad0c9..e885bc255aa 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp1_pp2_resume_torch_dist_reshard_2x1x4_te_8experts2parallel_dist_optimizer/model_config.yaml +++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp1_pp2_resume_torch_dist_reshard_2x1x4_te_8experts2parallel_dist_optimizer/model_config.yaml @@ -57,3 +57,4 @@ MODEL_ARGS: --no-bias-gelu-fusion: true --log-memory-to-tensorboard: true TEST_TYPE: ckpt-resume +LAUNCHER: ft_launcher diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_resume_torch_dist_te_8experts2parallel_multi_dist_optimizer_instances/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_resume_torch_dist_te_8experts2parallel_multi_dist_optimizer_instances/model_config.yaml index 159a7d1f4c2..37224a25cd2 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_resume_torch_dist_te_8experts2parallel_multi_dist_optimizer_instances/model_config.yaml +++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp1_resume_torch_dist_te_8experts2parallel_multi_dist_optimizer_instances/model_config.yaml @@ -63,3 +63,4 @@ MODEL_ARGS: --async-save: true --use-persistent-ckpt-worker: true TEST_TYPE: ckpt-resume +LAUNCHER: ft_launcher diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer/model_config.yaml index 7c292c0fa53..ff52550c238 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer/model_config.yaml +++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer/model_config.yaml @@ -67,3 +67,4 @@ MODEL_ARGS: --async-save: true --use-persistent-ckpt-worker: true TEST_TYPE: ckpt-resume +LAUNCHER: ft_launcher