From 994c6807703ba7f3e0790add4db03ba3415ba312 Mon Sep 17 00:00:00 2001 From: Deepak Narayanan Date: Sat, 24 Jan 2026 12:52:04 -0800 Subject: [PATCH] Keep FSDP's and DDP's finish_grad_sync API identical Signed-off-by: Deepak Narayanan --- .../core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py index c1c11721f7e..f3708a35dd8 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -1105,7 +1105,7 @@ def attach_grad_to_optimizer_state(self): """ self.param_and_grad_buffer.update_main_grads() - def finish_grad_sync(self): + def finish_grad_sync(self, force_all_reduce: Optional[bool] = False): """ Finishes grad sync (all-reduce or reduce-scatter) communication operations for all model gradients. Call prior to the optimization step to resolve @@ -1114,6 +1114,9 @@ def finish_grad_sync(self): When overlap_grad_reduce is set to True, waits for asynchronous communication calls to complete. When overlap_grad_reduce is set to False, calls synchronous communication ops. + + NOTE: force_all_reduce is included as an argument to maintain API compatibility + with DDP.force_grad_sync. """ # Synchronize gradient reduce-scatter operations for all model gradients. self.synchronize_gradient_reduce()