From 2c193add93b32deb48f2e39cd5375c01375249ce Mon Sep 17 00:00:00 2001 From: Wen Chen Date: Tue, 30 Jun 2026 19:07:54 +0000 Subject: [PATCH 1/2] examples/llama: disable gradient accumulation fusion on FSDP paths The fusion perf toggle defaults ON, but gradient accumulation fusion is incompatible with both FSDP backends: torch-FSDP2 hard-asserts against it in arguments.py (--use-torch-fsdp2 is not supported with gradient accumulation fusion), and Megatron-FSDP crashes at runtime with an fsdp_grads "No buffer found for bucket_id" assertion. Since the test harness drives FSDP=1, the default-on toggle broke every FSDP run. Force GRADIENT_ACCUMULATION_FUSION=0 inside the existing FSDP/MEGATRON_FSDP block (with a notice). The non-FSDP perf path keeps fusion on. Verified on 8xMI355X (MODEL_SIZE=8, mock data): FSDP=1 and MEGATRON_FSDP=1 both train, non-FSDP keeps fusion enabled. Co-Authored-By: Claude Opus 4.8 (1M context) --- examples/llama/train_llama2.sh | 7 +++++++ examples/llama/train_llama3.sh | 7 +++++++ 2 files changed, 14 insertions(+) diff --git a/examples/llama/train_llama2.sh b/examples/llama/train_llama2.sh index 126c266ece8..8bb8f8e92d2 100755 --- a/examples/llama/train_llama2.sh +++ b/examples/llama/train_llama2.sh @@ -123,6 +123,13 @@ fi if [ "$FSDP" -eq 1 ] || [ "$MEGATRON_FSDP" -eq 1 ]; then unset CUDA_DEVICE_MAX_CONNECTIONS + # Gradient accumulation fusion is incompatible with FSDP: torch-FSDP2 hard-asserts against it + # in arguments.py, and Megatron-FSDP crashes at runtime with an fsdp_grads "No buffer found for + # bucket_id" assertion. Force it off so the default-on toggle does not break the FSDP suites. + if [ "$GRADIENT_ACCUMULATION_FUSION" -eq 1 ]; then + echo "FSDP is incompatible with gradient accumulation fusion; disabling fusion (GRADIENT_ACCUMULATION_FUSION=0)." + GRADIENT_ACCUMULATION_FUSION=0 + fi if [ "$TP" -gt 1 ]; then echo "It is not recommended to use FSDP and TP together. Disabling TP." TP=1 diff --git a/examples/llama/train_llama3.sh b/examples/llama/train_llama3.sh index c5755f1c26f..edacb9f3ce4 100755 --- a/examples/llama/train_llama3.sh +++ b/examples/llama/train_llama3.sh @@ -137,6 +137,13 @@ fi if [ "$FSDP" -eq 1 ] || [ "$MEGATRON_FSDP" -eq 1 ]; then unset CUDA_DEVICE_MAX_CONNECTIONS + # Gradient accumulation fusion is incompatible with FSDP: torch-FSDP2 hard-asserts against it + # in arguments.py, and Megatron-FSDP crashes at runtime with an fsdp_grads "No buffer found for + # bucket_id" assertion. Force it off so the default-on toggle does not break the FSDP suites. + if [ "$GRADIENT_ACCUMULATION_FUSION" -eq 1 ]; then + echo "FSDP is incompatible with gradient accumulation fusion; disabling fusion (GRADIENT_ACCUMULATION_FUSION=0)." + GRADIENT_ACCUMULATION_FUSION=0 + fi if [ "$TP" -gt 1 ]; then echo "It is not recommended to use FSDP and TP together. Disabling TP." TP=1 From 7d8e7e44476017e7093a7847bbb871c23b95a183 Mon Sep 17 00:00:00 2001 From: Wen Chen Date: Tue, 30 Jun 2026 19:09:11 +0000 Subject: [PATCH 2/2] examples/llama: propagate training exit status from train scripts train_llama{2,3}.sh exited 0 even when torchrun crashed, so CI gating on exit code saw failures as passes. Two masks: (1) the run status was discarded, and with TEE_OUTPUT=1 the `... |& tee $TRAIN_LOG` pipeline returns tee's status (~always 0), not torchrun's; (2) the trailing throughput/mem perf-parse became the script's exit status. Wrap the run in `set -o pipefail` and capture $? (eval collapses PIPESTATUS to a single element = tee's, so PIPESTATUS[0] does not work here), emit an explicit error line on failure, and `exit $TRAIN_RC` at the end so the perf-parse still runs but cannot mask a crash. Verified: clean run exits 0; forced torchrun failure exits non-zero under both TEE_OUTPUT=0 and =1. Co-Authored-By: Claude Opus 4.8 (1M context) --- examples/llama/train_llama2.sh | 15 ++++++++++++++- examples/llama/train_llama3.sh | 15 ++++++++++++++- 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/examples/llama/train_llama2.sh b/examples/llama/train_llama2.sh index 8bb8f8e92d2..fa52559c729 100755 --- a/examples/llama/train_llama2.sh +++ b/examples/llama/train_llama2.sh @@ -473,8 +473,18 @@ else run_cmd="$run_cmd |& tee $TRAIN_LOG" fi -if [ "$NO_TRAINING" -eq 0 ]; then +TRAIN_RC=0 +if [ "$NO_TRAINING" -eq 0 ]; then + # pipefail so the run's `... |& tee $TRAIN_LOG` (TEE_OUTPUT=1) returns torchrun's status, not tee's + # (~always 0). Without it a training crash is masked and the script would exit 0. ($? after `eval` + # of the pipeline, since `eval` collapses PIPESTATUS to a single element = tee's status.) + set -o pipefail eval $run_cmd + TRAIN_RC=$? + set +o pipefail + if [ "$TRAIN_RC" -ne 0 ]; then + echo "ERROR: training (torchrun pretrain_gpt.py) failed with exit code $TRAIN_RC" |& tee -a $TRAIN_LOG + fi fi @@ -516,3 +526,6 @@ MEMUSAGE=$(python3 mean_log_value.py tmp.txt) echo "mem usages: $MEMUSAGE" |& tee -a "$TRAIN_LOG" rm tmp.txt +# Propagate the training exit status (the perf-parse above must not mask a torchrun failure). +exit $TRAIN_RC + diff --git a/examples/llama/train_llama3.sh b/examples/llama/train_llama3.sh index edacb9f3ce4..79a76b8165b 100755 --- a/examples/llama/train_llama3.sh +++ b/examples/llama/train_llama3.sh @@ -505,8 +505,18 @@ else run_cmd="$run_cmd |& tee $TRAIN_LOG" fi -if [ "$NO_TRAINING" -eq 0 ]; then +TRAIN_RC=0 +if [ "$NO_TRAINING" -eq 0 ]; then + # pipefail so the run's `... |& tee $TRAIN_LOG` (TEE_OUTPUT=1) returns torchrun's status, not tee's + # (~always 0). Without it a training crash is masked and the script would exit 0. ($? after `eval` + # of the pipeline, since `eval` collapses PIPESTATUS to a single element = tee's status.) + set -o pipefail eval $run_cmd + TRAIN_RC=$? + set +o pipefail + if [ "$TRAIN_RC" -ne 0 ]; then + echo "ERROR: training (torchrun pretrain_gpt.py) failed with exit code $TRAIN_RC" |& tee -a $TRAIN_LOG + fi fi @@ -547,3 +557,6 @@ grep -Eo 'mem usages: [^|]*' "$TRAIN_LOG" | sed -E 's/.*mem usages: ([0-9\.]+).* MEMUSAGE=$(python3 mean_log_value.py tmp.txt) echo "mem usages: $MEMUSAGE" |& tee -a "$TRAIN_LOG" rm tmp.txt + +# Propagate the training exit status (the perf-parse above must not mask a torchrun failure). +exit $TRAIN_RC