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 bd13e76379e..acc95805705 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -371,6 +371,11 @@ def _init_fsdp_param_and_grad_buffer(self): self.param_and_grad_buffer, ag_stream=self.side_stream_for_param_gather ) + # Also expose the reduce-scatter pipeline on the buffer so the lazy main_grad_getter can + # apply double-buffer back-pressure: it drains the oldest pending reduce-scatter before + # allocating a new grad bucket, preventing the fixed double-buffer pool from overflowing. + self.param_and_grad_buffer.grad_reduce_pipeline = self.grad_reduce_pipeline + # Set the suggested communication unit size for reduce-scatter and all-gather pipelines. suggested_communication_unit_size = self.ddp_config.suggested_communication_unit_size if suggested_communication_unit_size is None: diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py index 3ec117ebd9e..327f298278e 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py @@ -2398,10 +2398,20 @@ def _alloc(dtype, size): p._item_id = item_id def main_grad_getter(p): - # Make sure main_grad memory is allocated when initially accessed. - bucket = p._gbuf.fetch_bucket() gbuf = p._gbuf item_id = p._item_id + # Before lazily allocating this bucket's grad buffer, drain the oldest pending + # reduce-scatter if admitting it would push more than two FSDP units' grad + # buckets live at once. Without this back-pressure the fixed double-buffer pool + # (size 2) overflows during backward and FixedPoolAllocator.allocate asserts + # ("No buffer found for bucket_id"). Ports NVIDIA/Megatron-LM commit 55638bc4. + grad_reduce_pipeline = getattr(self, "grad_reduce_pipeline", None) + if grad_reduce_pipeline is not None: + grad_reduce_pipeline._enforce_double_buffer_limit( + [gbuf.bucket_index.bucket_id] + ) + # Make sure main_grad memory is allocated when initially accessed. + bucket = gbuf.fetch_bucket() # View it as p.shape so you can insert the param.grad into # the bucket seamlessly. return gbuf.get_item_from_bucket(bucket, item_id).view( @@ -3144,7 +3154,11 @@ def _enforce_double_buffer_limit(self, add_buckets): double_buf_units.add(fsdp_unit_id) if len(double_buf_units) > 2: keep_n -= 1 - self.wait_for_previous_grad_reduce(keep_n) + # Issue the drain (RS-completion waits + buffer frees) on the reduce-scatter stream rather + # than the caller's stream: this is also invoked lazily from main_grad_getter on the compute + # stream during backward, and the frees must order against the reduce-scatter, not compute. + with torch.cuda.stream(self.rs_stream): + self.wait_for_previous_grad_reduce(keep_n) def get_ready_bucket_group_for_reduction(self, bucket_id: int) -> Optional[List[int]]: """Checks if all buckets in the bucket group containing the given bucket_id