[Dev] Generalize FSDP persistent communication buffers

docs/user-guide/features/megatron_fsdp.md

@@ -344,7 +344,8 @@ Source: Feng, Wei, Will Constable, and Yifan Mao. “Getting Started with Fully
 | Optimization | Description | `Megatron-Core` Config | `fully_shard` Config |
 |--------------|-------------|----------------------|----------------------|
 | **FSDP Unit Modules** | A list of `str` or `class` import paths for `torch.nn.Module`(s) that are considered FSDP unit modules and sharded by Megatron-FSDP. Parameters and sub-modules that are not members of an FSDP unit are not sharded. |  Defaults to supported Megatron-Core modules (`TransformerLayer`, etc.) in Megatron-LM. | `fsdp_unit_modules=[...]` |
-| **FSDP Double Buffer Allocator** | Megatron-FSDP uses the double-buffer allocator, which persistently allocates a buffer pair assigned to alternating FSDP units that temporarily stores parameters and gradients. Automatically used with NCCL user buffer registration. | `--fsdp-double-buffer` | `fsdp_double_buffer=True` |
+| **FSDP Persistent Buffer Pool** | Enables persistent communication-buffer pools. The legacy `double_buffer` name is retained for compatibility, but pool capacity is configurable and defaults to two. Automatically used with NCCL user buffer registration. | `--fsdp-double-buffer` | `fsdp_double_buffer=True` |
+| **FSDP Persistent Buffer Count** | Number of buffers in each persistent communication pool. Must be at least two when the pool is enabled; changing the default while the pool is disabled is invalid. | `--fsdp-buffer-count <int>` | `fsdp_buffer_count=<int>` |
 | **FSDP Max Pool Allocator** | Megatron-FSDP uses the `MaxPoolAllocator`, which supports double buffering hybrid / asymmetrical model architectures by taking the maximum of all layers. Automatically sets `--fsdp-double-buffer`. | `--megatron-fsdp-max-pool-double-buffer` | `maxpool_double_buffer=True` |
 | **Param All-Gather Overlap** | Whether to overlap parameter all-gather with compute. Automatically activated for the ZeRO-3 sharding strategy. | `--overlap-param-gather` | `overlap_param_gather=True` |
 | **Gradient Reduce-Scatter Overlap** | Whether to overlap gradient reduce-scatter or all-reduce with compute. Automatically activated for ZeRO-2 and ZeRO-3 sharding strategies. | `--overlap-grad-reduce` | `overlap_grad_reduce=True` |
@@ -395,28 +396,28 @@ To implement these "unit-periodic" mechanics, Megatron-FSDP uses `Module` hooks
   - When `module.state_dict()` (for any module managed by Megatron-FSDP) is invoked, Megatron-FSDP will swap all parameter references to point to sharded `DTensor` main weights for distributed optimization and checkpointing.
   - When `MegatronFSDP.load_state_dict()` is invoked, both the main and compute weights are updated. When using quantized model compute, the main weights are quantized and sharded.
 
-#### Double Buffering
+#### Persistent Buffer Pools
 
 Megatron-FSDP uses a `Tensor._typed_storage()._resize_(bytes)`-based allocator to instantly allocate and de-allocate memory without depending on the `CUDACachingAllocator` for un-sharded parameters and gradients by default. (Cache fragmentation and garbage collection can procrastinate large quantities of `cudaMalloc` and `cudaFree` operations that can block programs and spike memory, particularly when memory utilization is maxed out.) However, modifying the underlying storage of a buffer is not compatible with NCCL symmetric registration or CUDA graphability, which require a persistent state during runtime.
 
-To support these optimizations, Megatron-FSDP uses **double-buffering**, which assigns 2 persistently-allocated buffers to FSDP units in an alternating pattern, hard-limiting the memory overhead for parameter and gradient buffer allocation and ensuring that no more than 2 FSDP units are computed or communicated concurrently.
+To support these optimizations, Megatron-FSDP can use persistent communication-buffer pools. The legacy `fsdp_double_buffer` option enables these pools, while `fsdp_buffer_count` controls the capacity of each pool. Its default value of 2 preserves conventional double-buffer behavior, but larger values can support schedules with more concurrently live FSDP units.
 
 ```{figure} ../../images/megatron_fsdp/fsdp_double_buffer.png
 :alt: FSDP Double Buffering
 :align: center
 
-Visualization of double buffering in Megatron-FSDP. Even- and odd-indexed FSDP units share the same un-sharded parameter and gradient buffers, overwriting incumbent data as needed during runtime. Megatron-FSDP ensures that no more than two FSDP units are un-sharded at any point during runtime.
+Visualization of the default two-slot persistent pool in Megatron-FSDP. Slots are shared capacity assigned to compatible live buckets at runtime rather than statically dedicated to pipeline-schedule roles. With `fsdp_buffer_count=3`, for example, a current forward unit and two forward-prefetched successors may occupy three slots when the schedule permits it.
 ```
 
-With double-buffering, Megatron-FSDP does not need to allocate memory after initialization, which can reduce memory fragmentation and improve performance. However, double-buffering requires _depth-wise model symmetry_, where even- and odd-indexed FSDP units have identical size during runtime. If double-buffering is utilized, Megatron-FSDP computes the **_mode_** of FSDP unit sizes as the symmetrical double-buffer size, and any FSDP units not symmetrical to the computed size will default to the `_resize_(bytes)`-based allocator (or persistently allocated for extremely large and asymmetrical layers that affect performance significantly like `torch.nn.Embedding` when the low-level argument `fsdp_db_use_persist_buf_on_alloc_fail` is set).
+With persistent pools, Megatron-FSDP does not need to allocate memory after initialization, which can reduce memory fragmentation and improve performance. The `FixedPoolAllocator` requires _depth-wise model symmetry_, where FSDP units have identical bucket layouts during runtime. Megatron-FSDP computes the **_mode_** of FSDP unit sizes as the fixed-pool size, and any FSDP units not symmetrical to the computed size will default to the `_resize_(bytes)`-based allocator (or be persistently allocated for extremely large and asymmetrical layers that affect performance significantly, such as `torch.nn.Embedding`, when the low-level argument `fsdp_db_use_persist_buf_on_alloc_fail` is set).
 
-Not all model architectures support depth-wise model symmetry. For example, hybrid architectures like **Nemotron** are a combination of Transformer, Mamba, and MoE blocks that are asymmetrical in size and data-type. To double-buffer these model architectures, we need a pool of buffers that can support any FSDP unit, which can be computed from the _**maximum**_ of all FSDP units, and this "MaxPool" of (now symmetric) buffers of maximum size, shape, and dtype can be double-buffered.
+Not all model architectures support depth-wise model symmetry. For example, hybrid architectures like **Nemotron** are a combination of Transformer, Mamba, and MoE blocks that are asymmetrical in size and data-type. To persistently pool buffers for these model architectures, Megatron-FSDP can compute a pool from the _**maximum**_ requirements of all FSDP units. This `MaxPool` supports the configured number of concurrently live units even when their layouts differ.
 
 ```{figure} ../../images/megatron_fsdp/maxpool_allocator.png
 :alt: MaxPoolAllocator
 :align: center
 
-Visualizing the MaxPoolAllocator initialization in Megatron-FSDP. Iterating through all FSDP units, data buckets are categorized by data-type, sorted from small to large, and compared to the current MaxPool. If there are not enough buckets in the pool to support the unit, buckets are added to the pool (with size 0). If the largest buckets of the pool are not large enough to support the buckets in the unit (assigned to the pool from smallest to largest), the buckets in the pool are enlarged. After this process, we arrive at a minimal set of buckets that can double-buffer every FSDP unit in the model.
+Visualizing the MaxPoolAllocator initialization in Megatron-FSDP. Iterating through all FSDP units, data buckets are categorized by data-type, sorted from small to large, and compared to the current MaxPool. If there are not enough buckets in the pool to support the unit, buckets are added to the pool (with size 0). If the largest buckets of the pool are not large enough to support the buckets in the unit (assigned to the pool from smallest to largest), the buckets in the pool are enlarged. After this process, we arrive at a minimal bucket layout that can persistently buffer every FSDP unit in the model, replicated according to `fsdp_buffer_count`.
 ```
 
 ### Data-Parallel Sharding Strategies
@@ -600,7 +601,7 @@ Megatron-FSDP sharding and communication buffers support mixed-precision, such t
 
 | Optimization | Description | `Megatron-Core` Config | `fully_shard` Config |
 |--------------|-------------|----------------------|----------------------|
-| **NCCL User Buffers** | Allocate and register Megatron-FSDP communication buffers with NCCL, which enables zero-`COPY`, high-precision reduction, copy-engine collectives, and symmetric kernels. Uses double buffering. | `--use-nccl-ub` | `nccl_ub=True` |
+| **NCCL User Buffers** | Allocate and register Megatron-FSDP communication buffers with NCCL, which enables zero-`COPY`, high-precision reduction, copy-engine collectives, and symmetric kernels. Automatically enables persistent buffer pools with a default capacity of two. | `--use-nccl-ub` | `nccl_ub=True` |
 | **NCCL Manual Registration** | Instead of registering NCCL user buffers on first allocation, batch registration of all communication buffers at the end of the initial training step. Reduces registration latency. | `--fsdp-manual-registration` | N/A (Megatron-Core Only) |
 | **Disable Symmetric Registration** | Disable symmetric registration with NCCL. Optional, as symmetric registration failure defaults to normal registration. | `--disable-symmetric-registration` | `disable_symmetric_registration=True` |
 

megatron/core/distributed/distributed_data_parallel_config.py

@@ -145,11 +145,19 @@ class DistributedDataParallelConfig:
     """
 
     fsdp_double_buffer: bool = False
-    """If true, use persistently allocated double buffers for the 
-      temporary memory needed in the Megatron FSDP communications.
-      This option will cause additional memory overhead, however, it is necessary for
-      to register user buffer (nccl_ub=True) for the Megatron FSDP. 
-      This option will be automatically set to True when nccl_ub=True.
+    """If true, enable persistent communication-buffer pools for Megatron FSDP.
+      The legacy option name is retained for compatibility; the number of buffers in
+      each pool is controlled by fsdp_buffer_count. Persistent buffers add memory
+      overhead but are required for NCCL user-buffer registration. This option is
+      automatically enabled when nccl_ub=True or the MaxPool allocator is requested.
+    """
+
+    fsdp_buffer_count: int = 2
+    """Number of persistent buffers allocated for each Megatron FSDP communication pool.
+      The default of two preserves the conventional double-buffer behavior. Combined 1F1B
+      overlap may require three buffers because a backward/recompute unit, the current
+      forward unit, and its forward-prefetched successor can be live concurrently. This
+      option may only be changed when fsdp_double_buffer is enabled.
     """
 
     fsdp_db_use_persist_buf_on_alloc_fail: bool = False
@@ -306,6 +314,17 @@ def __post_init__(self):
             assert self.bucket_size is None, "Cannot specify both num_buckets and bucket_size"
             assert self.num_buckets > 0, "num_buckets must be greater than 0"
 
-        if self.megatron_fsdp_max_pool_double_buffer:
-            # MaxPoolAllocator is a type of double-buffer allocator.
+        if self.nccl_ub or self.megatron_fsdp_max_pool_double_buffer:
+            # NCCL user buffers and MaxPoolAllocator require persistent buffer pools.
             self.fsdp_double_buffer = True
+
+        if self.fsdp_double_buffer:
+            if self.fsdp_buffer_count < 2:
+                raise ValueError(
+                    "fsdp_buffer_count must be at least 2 when fsdp_double_buffer is enabled."
+                )
+        elif self.fsdp_buffer_count != 2:
+            raise ValueError(
+                "fsdp_buffer_count may only be changed from its default of 2 when "
+                "fsdp_double_buffer is enabled."
+            )

megatron/core/distributed/fsdp/mcore_fsdp_adapter.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -167,10 +167,13 @@ def __init__(
         self._annotate_tensor_parallelism(module)
 
         if config.overlap_moe_expert_parallel_comm:
-            assert not ddp_config.fsdp_double_buffer, (
-                "1F1B overlap with FSDP does not support double buffer. "
-                "Please set fsdp_double_buffer=False in the ddp config."
-            )
+            if ddp_config.fsdp_double_buffer:
+                assert ddp_config.fsdp_buffer_count >= 3, (
+                    "1F1B overlap with persistent Megatron-FSDP communication buffers "
+                    "requires fsdp_buffer_count >= 3. A backward/recompute unit, the "
+                    "current forward unit, and its forward-prefetched successor can be "
+                    "live concurrently."
+                )
             assert config.cuda_graph_impl in ("none", "full_iteration"), (
                 "1F1B overlap with FSDP does not support per-layer CUDA graphs "
                 f"(cuda_graph_impl={config.cuda_graph_impl!r}). "

megatron/core/distributed/fsdp/src/README.md

@@ -160,8 +160,11 @@ Megatron-FSDP's `fully_shard_*` API has a comprehensive set of arguments for fin
     - Defaults to `False`, but will be automatically enabled in Megatron-LM.
 - `disable_symmetric_registration` will disable NCCL window (i.e. symmetric) registration when using `nccl_ub`. 
     - Defaults to `False`.
-- `fsdp_double_buffer` will use persistently allocated double buffers for temporarily-defined memory needed in `MegatronFSDP` communications. Having persistent double buffers may increase peak VRAM utilization, but is required to register NCCL user buffers (`nccl_ub=True`) for `MegatronFSDP`. Currently, this is only supported for simple repetitive model structures such as GPT.
-    - Defaults to `False`. Automatically overridden to `True` when `nccl_ub` is enabled.
+- `fsdp_double_buffer` is the backward-compatible enable flag for persistent communication-buffer pools used by `MegatronFSDP`. Despite the legacy name, it does not require the pool capacity to remain at two. Having persistent buffers may increase peak VRAM utilization, but is required to register NCCL user buffers (`nccl_ub=True`) for `MegatronFSDP`.
+    - Defaults to `False`. Automatically overridden to `True` when `nccl_ub` or `maxpool_double_buffer` is enabled.
+- `fsdp_buffer_count` controls the number of persistent buffers in each FSDP communication pool.
+    - Defaults to two. A non-default value is invalid unless the persistent pool is enabled. Combined 1F1B overlap with forward prefetch requires at least three because a backward/recompute unit, the current forward unit, and its prefetched successor can be live concurrently.
+    - Slots are shared pool capacity assigned to compatible live buckets at runtime; they are not statically dedicated to particular pipeline-schedule roles.
 - `maxpool_double_buffer` will use a max-pooling algorithm to build a sufficient pool of buffers that can support all layers of hybrid / asymmetrical model architectures like Nemotron.
     - Defaults to `False`. Highly-recommended for hybrid architectures when using `fsdp_double_buffer=True` to double-buffer every layer of the model.
 - `preproc_state_dict_for_dcp_ckpt` adds `model.state_dict()` and `optimizer.state_dict()` post-hooks that modify the model and optimizer state in preparation for `torch.distributed.checkpoint.{save,load}` ([Torch DCP](https://docs.pytorch.org/docs/stable/distributed.checkpoint.html)) checkpointing. Specifically, it adds `__create_write_items__` and `__create_chunk_list__` methods to Tensors utilized by Torch DCP to redistribute parameters when saving and loading model and optimizer checkpoints. Can be deactivated should the user need a custom distributed checkpointing strategy.

megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py

@@ -76,11 +76,19 @@ class DistributedDataParallelConfig:
     """
 
     fsdp_double_buffer: bool = False
-    """If true, use persistently allocated double buffers for the 
-      temporary memory needed in the Megatron FSDP communications.
-      This option will cause additional memory overhead, however, it is necessary for
-      to register user buffer (nccl_ub=True) for the Megatron FSDP. 
-      This option will be automatically set to True when nccl_ub=True.
+    """If true, enable persistent communication-buffer pools for Megatron FSDP.
+      The legacy option name is retained for compatibility; the number of buffers in
+      each pool is controlled by fsdp_buffer_count. Persistent buffers add memory
+      overhead but are required for NCCL user-buffer registration. This option is
+      automatically enabled when nccl_ub=True or the MaxPool allocator is requested.
+    """
+
+    fsdp_buffer_count: int = 2
+    """Number of persistent buffers allocated for each Megatron FSDP communication pool.
+      The default of two preserves the conventional double-buffer behavior. Combined 1F1B
+      overlap may require three buffers because a backward/recompute unit, the current
+      forward unit, and its forward-prefetched successor can be live concurrently. This
+      option may only be changed when fsdp_double_buffer is enabled.
     """
 
     fsdp_all_gather_in_start_param_sync: bool = True
@@ -213,6 +221,17 @@ def __post_init__(self):
                     "with nccl_ub due to compatibility issue with torch.cuda.MemPool API."
                 )
 
-        if self.megatron_fsdp_max_pool_double_buffer:
-            # MaxPoolAllocator is a type of double-buffer allocator.
+        if self.nccl_ub or self.megatron_fsdp_max_pool_double_buffer:
+            # NCCL user buffers and MaxPoolAllocator require persistent buffer pools.
             self.fsdp_double_buffer = True
+
+        if self.fsdp_double_buffer:
+            if self.fsdp_buffer_count < 2:
+                raise ValueError(
+                    "fsdp_buffer_count must be at least 2 when fsdp_double_buffer is enabled."
+                )
+        elif self.fsdp_buffer_count != 2:
+            raise ValueError(
+                "fsdp_buffer_count may only be changed from its default of 2 when "
+                "fsdp_double_buffer is enabled."
+            )