Summary. The training loop wraps each micro-step in distributed.ddp_sync_grad(model, is_boundary) to defer gradient synchronization to the last micro-step of an accumulation window (cosmos_framework/trainer/__init__.py, training_step). The context manager only handles DistributedDataParallel instances (cosmos_framework/utils/distributed.py). With trainer.distributed_parallelism="fsdp" the model is a set of fully_shard (FSDP2) modules, the isinstance check fails, and the context silently does nothing. Every micro-step then runs the full gradient reduction. Accumulation enlarges the batch but saves no communication.
Measurement (2x 8xA100-80GB, Ethernet inter-node, HSDP shard=8 replicate=2, action-policy recipe, batch 64/rank):
| Config |
s/iter |
samples/s |
| grad_accum_iter=1 |
110.2 |
9.3 |
| grad_accum_iter=4 |
430.6 |
9.5 |
430.6 s is four accum=1 steps back to back.
Fix, validated on hardware (same instance type, a different rental whose accum=1 step is 22.2 s): defer only the cross-replica all-reduce via FSDPModule.set_requires_all_reduce(False) on non-boundary micro-steps. The intra-node reduce-scatter still runs each micro-step, so gradients stay sharded and per-rank memory does not grow. With the fix, accum=4 steps take 54.3 / 54.3 / 54.9 s against 88.8 s (4 x 22.2) for the unamortized path. The two measurements decompose to compute 10.8 s + sync 11.4 s per micro-batch, throughput 1.63x at accum=4. set_requires_gradient_sync(False) also works but keeps gradients unsharded across the window, which adds the full unsharded gradient bytes per rank.
PR with the fix follows from Sullivan07043:fix-fsdp2-accum.
Environment. cosmos-framework main (ee58e41), torch 2.10 cu128, NCCL 2.27.5, 16 ranks over 2 nodes.
Summary. The training loop wraps each micro-step in
distributed.ddp_sync_grad(model, is_boundary)to defer gradient synchronization to the last micro-step of an accumulation window (cosmos_framework/trainer/__init__.py,training_step). The context manager only handlesDistributedDataParallelinstances (cosmos_framework/utils/distributed.py). Withtrainer.distributed_parallelism="fsdp"the model is a set offully_shard(FSDP2) modules, the isinstance check fails, and the context silently does nothing. Every micro-step then runs the full gradient reduction. Accumulation enlarges the batch but saves no communication.Measurement (2x 8xA100-80GB, Ethernet inter-node, HSDP shard=8 replicate=2, action-policy recipe, batch 64/rank):
430.6 s is four accum=1 steps back to back.
Fix, validated on hardware (same instance type, a different rental whose accum=1 step is 22.2 s): defer only the cross-replica all-reduce via
FSDPModule.set_requires_all_reduce(False)on non-boundary micro-steps. The intra-node reduce-scatter still runs each micro-step, so gradients stay sharded and per-rank memory does not grow. With the fix, accum=4 steps take 54.3 / 54.3 / 54.9 s against 88.8 s (4 x 22.2) for the unamortized path. The two measurements decompose to compute 10.8 s + sync 11.4 s per micro-batch, throughput 1.63x at accum=4.set_requires_gradient_sync(False)also works but keeps gradients unsharded across the window, which adds the full unsharded gradient bytes per rank.PR with the fix follows from
Sullivan07043:fix-fsdp2-accum.Environment. cosmos-framework main (ee58e41), torch 2.10 cu128, NCCL 2.27.5, 16 ranks over 2 nodes.