From 555c09f15954ac2331011632dd20abd30a8c2a74 Mon Sep 17 00:00:00 2001 From: Sullivan07043 Date: Tue, 4 Aug 2026 17:50:20 -0700 Subject: [PATCH 1/2] Fix gradient accumulation under FSDP2: defer cross-replica all-reduce ddp_sync_grad only handled DistributedDataParallel. Under FSDP2 (fully_shard) it was a silent no-op, so every micro-step of an accumulation window ran the full cross-node gradient reduction and grad_accum_iter saved no communication (measured: accum=4 step = 4x the accum=1 step on 2x8 A100 over Ethernet). Defer only the cross-replica all-reduce via FSDPModule.set_requires_all_reduce on non-boundary micro-steps. The intra-node reduce-scatter still runs every micro-step, so gradients stay sharded and per-rank memory does not grow. Validated on 2x 8xA100-80GB (HSDP shard=8 replicate=2): accum=4 step 54.5 s vs 88.8 s unamortized on a fleet whose accum=1 step is 22.2 s. --- cosmos_framework/utils/distributed.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/cosmos_framework/utils/distributed.py b/cosmos_framework/utils/distributed.py index 87cdd39e..1c6d5bf0 100644 --- a/cosmos_framework/utils/distributed.py +++ b/cosmos_framework/utils/distributed.py @@ -241,6 +241,7 @@ def ddp_sync_grad(model, enabled): else gradients will still be synchronized. """ assert isinstance(model, torch.nn.Module) + fsdp_modules: list = [] if isinstance(model, DistributedDataParallel): old_require_backward_grad_sync = model.require_backward_grad_sync if model.static_graph and model.require_backward_grad_sync != enabled: @@ -249,11 +250,28 @@ def ddp_sync_grad(model, enabled): model.show_sync_grad_static_graph_warning = False else: model.require_backward_grad_sync = enabled + else: + # FSDP2 (fully_shard) branch. Without this, gradient accumulation + # reduces gradients on EVERY micro-step and saves no communication + # (measured: accum=4 step = 4x the accum=1 step on Ethernet HSDP). + # We defer only the cross-replica all-reduce to the boundary + # micro-step. The intra-node reduce-scatter still runs every + # micro-step, so gradients stay sharded and VRAM does not grow + # (set_requires_gradient_sync(False) would keep gradients + # unsharded: +stored-grad-bytes per rank). + from torch.distributed.fsdp import FSDPModule + + fsdp_modules = [m for m in model.modules() if isinstance(m, FSDPModule)] + for m in fsdp_modules: + m.set_requires_all_reduce(enabled, recurse=False) try: yield finally: if isinstance(model, DistributedDataParallel): model.require_backward_grad_sync = old_require_backward_grad_sync + else: + for m in fsdp_modules: + m.set_requires_all_reduce(True, recurse=False) def collate_batches(data_batches: list[dict[str, torch.Tensor]]) -> torch.Tensor | dict[str, torch.Tensor]: From 799dbbabb2870db15631c58d2ff9aba0c5747a42 Mon Sep 17 00:00:00 2001 From: Sullivan07043 Date: Thu, 6 Aug 2026 16:24:38 -0700 Subject: [PATCH 2/2] Address review: restore previous all-reduce state, mutate inside try/finally - Capture each FSDP module's prior all_reduce_grads (mirroring the setter's write target) and restore exactly that value on exit, so nested or external sync contexts are preserved. - Record (module, prev) as mutations happen inside the protected block; if the setup loop raises midway, the finally block restores exactly the modules already changed. - Same treatment for the DDP branch via a mutation flag. Verified with a single-rank fully_shard test: flip/restore, nested contexts preserving outer state, exception-path restore, enabled=True no-op. --- cosmos_framework/utils/distributed.py | 67 ++++++++++++++++----------- 1 file changed, 40 insertions(+), 27 deletions(-) diff --git a/cosmos_framework/utils/distributed.py b/cosmos_framework/utils/distributed.py index 1c6d5bf0..bdc0e03e 100644 --- a/cosmos_framework/utils/distributed.py +++ b/cosmos_framework/utils/distributed.py @@ -241,37 +241,50 @@ def ddp_sync_grad(model, enabled): else gradients will still be synchronized. """ assert isinstance(model, torch.nn.Module) - fsdp_modules: list = [] - if isinstance(model, DistributedDataParallel): - old_require_backward_grad_sync = model.require_backward_grad_sync - if model.static_graph and model.require_backward_grad_sync != enabled: - if model.show_sync_grad_static_graph_warning: - log.warning("DDP static_graph=True is incompatible with sync_grad(). Performance will be reduced.") - model.show_sync_grad_static_graph_warning = False - else: - model.require_backward_grad_sync = enabled - else: - # FSDP2 (fully_shard) branch. Without this, gradient accumulation - # reduces gradients on EVERY micro-step and saves no communication - # (measured: accum=4 step = 4x the accum=1 step on Ethernet HSDP). - # We defer only the cross-replica all-reduce to the boundary - # micro-step. The intra-node reduce-scatter still runs every - # micro-step, so gradients stay sharded and VRAM does not grow - # (set_requires_gradient_sync(False) would keep gradients - # unsharded: +stored-grad-bytes per rank). - from torch.distributed.fsdp import FSDPModule - - fsdp_modules = [m for m in model.modules() if isinstance(m, FSDPModule)] - for m in fsdp_modules: - m.set_requires_all_reduce(enabled, recurse=False) + ddp_mutated = False + old_require_backward_grad_sync = None + fsdp_prev: list = [] # (FSDPModule, previous all_reduce_grads) for exactly the modules we mutated try: + if isinstance(model, DistributedDataParallel): + old_require_backward_grad_sync = model.require_backward_grad_sync + if model.static_graph and model.require_backward_grad_sync != enabled: + if model.show_sync_grad_static_graph_warning: + log.warning("DDP static_graph=True is incompatible with sync_grad(). Performance will be reduced.") + model.show_sync_grad_static_graph_warning = False + else: + model.require_backward_grad_sync = enabled + ddp_mutated = True + else: + # FSDP2 (fully_shard) branch. Without this, gradient accumulation + # reduces gradients on EVERY micro-step and saves no communication + # (measured: accum=4 step = 4x the accum=1 step on Ethernet HSDP). + # We defer only the cross-replica all-reduce to the boundary + # micro-step. The intra-node reduce-scatter still runs every + # micro-step, so gradients stay sharded and VRAM does not grow + # (set_requires_gradient_sync(False) would keep gradients + # unsharded: +stored-grad-bytes per rank). + from torch.distributed.fsdp import FSDPModule + + for m in model.modules(): + if not isinstance(m, FSDPModule): + continue + # Mirror of set_requires_all_reduce(recurse=False): it writes + # state._fsdp_param_group.all_reduce_grads when the group is + # truthy, so that is the previous state to capture. Recording + # (module, prev) as we mutate keeps the finally-block exact + # even if this loop raises midway. + fsdp_param_group = m._get_fsdp_state()._fsdp_param_group + if not fsdp_param_group: + continue + prev = fsdp_param_group.all_reduce_grads + m.set_requires_all_reduce(enabled, recurse=False) + fsdp_prev.append((m, prev)) yield finally: - if isinstance(model, DistributedDataParallel): + if ddp_mutated: model.require_backward_grad_sync = old_require_backward_grad_sync - else: - for m in fsdp_modules: - m.set_requires_all_reduce(True, recurse=False) + for m, prev in fsdp_prev: + m.set_requires_all_reduce(prev, recurse=False) def collate_batches(data_batches: list[dict[str, torch.Tensor]]) -> torch.Tensor | dict[str, torch.Tensor]: