examples/llama: fix FSDP grad-accum-fusion regression + propagate train exit status#145
Conversation
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
|
Megatron-FSDP should support Gradient accumulation fusion. The error that you saw was a bug from the upstream which was fixed recently in the commit: https://github.com/NVIDIA/Megatron-LM/pull/5222/changes Basically, the error occurs only when you use double buffers for ping-ponging, when you disable it, it runs fine for me with gradient_accumulation_fusion. [2026-06-30 21:40:55.941706] iteration 9/ 12 | consumed samples: 2304 | elapsed time per iteration (ms): 32108.9 | mem usages: 0.6143 | throughput per GPU (TFLOP/s/GPU): 420.2 | learning rate: 1.000000E-04 | global batch size: 256 | lm loss: 3.268351E-02 | loss scale: 1.0 | grad norm: 0.202 | num zeros: 0 | number of skipped iterations: 0 | number of nan iterations: 0 | The issue happens because of overflow of requirement for inflight buffers when we only use 2 buffers for AG and RS. For now, I think it's fine to disable MegatronFSDP when used with gradient accumulation fusion, I can enable this path again in my PR. |
|
@sudhu2k Let's get it disabled for now to unblock. Then once the fix is in, we can reenable it. |
2 participants
Summary
Two independent fixes to
examples/llama/train_llama{2,3}.sh:Disable gradient accumulation fusion on FSDP paths. The perf toggle added in examples/llama: default-on fusion & DP-comm perf toggles #141 defaults
GRADIENT_ACCUMULATION_FUSIONon, but fusion is incompatible with both FSDP backends, so every FSDP run (the test harness drives--fsdp 1) broke:FSDP=1FR--use-torch-fsdp2) hard-asserts inarguments.py:--use-torch-fsdp2 is not suu pported with gradient accumulation fusion.MEGATRON_FSDP=1) crashes at runtime:[FSDP][Rank N][fsdp_grads] No buffer found for bucket_id ....Fix: force
GRADIENT_ACCUMULATION_FUSION=0inside the existingFSDP || MEGATRON_FSDPblock (with a notice). The non-FSDP path keeps fusion on.Propagate the training exit status. The scripts exited
0even whentorchruncrashed, so CI gating on exit code saw failures as passes ~@~T which is how the examples/llama: default-on fusion & DP-comm perf toggles #141 regression slipped through green. Two masks:eval $run_cmdsttatus was discarded (and with
TEE_OUTPUT=1the... |& teepipeline returnstee's status, ~0), and the trailing perf-parse became the script's exit code. Fix:set -o pipefail+ capture$?, emit an explicit error line, andexit $TRAIN_RCat the end.Verification (8~WMI355X, MODEL_SIZE=8, mock data, 3 iters)
Immediate unblock without the script change: run with
GRADIENT_ACCUMULATION_FUSION=0.