[Dev] Add Megatron-FSDP weight prefetch for full recompute

megatron/core/distributed/distributed_data_parallel_config.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 from dataclasses import dataclass
 from typing import Optional, Tuple
@@ -215,6 +215,16 @@ class DistributedDataParallelConfig:
       main gradients to parameter dtype for `.grad`.
     """
 
+    megatron_fsdp_prefetch_recompute_forward_weights: bool = False
+    """If set to True, Megatron-FSDP prefetches rowwise weights needed by activation
+      recomputation during backward before prefetching backward transpose weights.
+    """
+
+    megatron_fsdp_cache_param_bucket_views: bool = False
+    """If set to True, Megatron-FSDP caches parameter bucket views to reduce repeated
+      Python-side view setup when attaching module parameters to all-gather buckets.
+    """
+
     megatron_fsdp_cuda_graph_mode: bool = False
     """If set to True, Megatron-FSDP will practice CUDA graph-safe operations, such as
     not dereferencing `param.grad` after the optimizer step to preserve references for
@@ -244,6 +254,20 @@ def __post_init__(self):
         if self.reuse_grad_buf_for_mxfp8_param_ag:
             assert self.fp8_param_gather, "Reuse grad buffer only when keeping params in MXFP8."
 
+        if self.megatron_fsdp_prefetch_recompute_forward_weights:
+            assert (
+                self.use_megatron_fsdp
+            ), "megatron_fsdp_prefetch_recompute_forward_weights requires use_megatron_fsdp."
+            assert self.data_parallel_sharding_strategy == "optim_grads_params", (
+                "megatron_fsdp_prefetch_recompute_forward_weights is only supported with "
+                "data_parallel_sharding_strategy='optim_grads_params'."
+            )
+
+        if self.megatron_fsdp_cache_param_bucket_views:
+            assert (
+                self.use_megatron_fsdp
+            ), "megatron_fsdp_cache_param_bucket_views requires use_megatron_fsdp."
+
         if self.nccl_ub and not is_torch_min_version("2.11.0a0"):
             if 'expandable_segments:True' in os.getenv('PYTORCH_CUDA_ALLOC_CONF', '').split(','):
                 raise ValueError(

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

@@ -1,4 +1,4 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 from dataclasses import dataclass
 from typing import Optional
@@ -155,6 +155,16 @@ class DistributedDataParallelConfig:
       main gradients to parameter dtype for `.grad`.
     """
 
+    megatron_fsdp_prefetch_recompute_forward_weights: bool = False
+    """If set to True, Megatron-FSDP prefetches rowwise weights needed by activation
+      recomputation during backward before prefetching backward transpose weights.
+    """
+
+    megatron_fsdp_cache_param_bucket_views: bool = False
+    """If set to True, Megatron-FSDP caches parameter bucket views to reduce repeated
+      Python-side view setup when attaching module parameters to all-gather buckets.
+    """
+
     megatron_fsdp_cuda_graph_mode: bool = False
     """If set to True, Megatron-FSDP will practice CUDA graph-safe operations, such as
     not dereferencing `param.grad` after the optimizer step to preserve references for
@@ -181,6 +191,12 @@ def __post_init__(self):
         import os
 
         """Check the validity of the config."""
+        if self.megatron_fsdp_prefetch_recompute_forward_weights:
+            assert self.data_parallel_sharding_strategy == "optim_grads_params", (
+                "megatron_fsdp_prefetch_recompute_forward_weights is only supported with "
+                "data_parallel_sharding_strategy='optim_grads_params'."
+            )
+
         if self.nccl_ub and not is_torch_min_version("2.11.0a0"):
             if 'expandable_segments:True' in os.getenv('PYTORCH_CUDA_ALLOC_CONF', '').split(','):
                 raise ValueError(

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

@@ -1,16 +1,4 @@
-# Copyright (c) 2025, NVIDIA CORPORATION.  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.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import logging
 import types
@@ -103,6 +91,8 @@ def fully_shard_model(
     fsdp_db_use_persist_buf_on_alloc_fail: bool = False,
     disable_symmetric_registration: bool = False,
     enable_fine_grained_param_gather: bool = False,
+    prefetch_recompute_forward_weights: bool = False,
+    cache_param_bucket_views: bool = False,
     use_decoupled_grad: bool = False,
     cuda_graph_mode: bool = False,
 ) -> torch.nn.Module:
@@ -262,6 +252,14 @@ class that schedules the sharding lifecycle of the model parameters and gradient
             unshards parameters per-Module instead of unsharding all sub-modules of an FSDP
             unit module simultaneously. Defaults to False.
 
+        prefetch_recompute_forward_weights (bool):
+            Whether to prefetch rowwise weights needed by activation recomputation during
+            backward before prefetching backward transpose weights. Defaults to False.
+
+        cache_param_bucket_views (bool):
+            Whether to cache parameter bucket views to reduce repeated Python-side view setup
+            when attaching module parameters to all-gather buckets. Defaults to False.
+
         use_decoupled_grad (bool):
             If true, reduced gradients are installed into `Parameter.decoupled_grad` instead
             of `Parameter.grad`. Defaults to False.
@@ -358,6 +356,17 @@ class that schedules the sharding lifecycle of the model parameters and gradient
             "Meta device initialization (init_model_with_meta_device=True) is not "
             "supported or necessary for the 'no_shard' / 0 sharding strategy."
         )
+    if prefetch_recompute_forward_weights:
+        if zero_dp_strategy != "optim_grads_params":
+            raise ValueError(
+                "prefetch_recompute_forward_weights is only supported with "
+                "zero_dp_strategy='optim_grads_params'."
+            )
+        if not fsdp_unit_modules:
+            raise ValueError(
+                "prefetch_recompute_forward_weights requires fsdp_unit_modules to define "
+                "the Megatron-FSDP unit-level backward prefetch order."
+            )
 
     # DDP Config for Megatron FSDP.
     ddp_config = DistributedDataParallelConfig(
@@ -371,6 +380,8 @@ class that schedules the sharding lifecycle of the model parameters and gradient
         fsdp_double_buffer=fsdp_double_buffer or nccl_ub,
         fsdp_db_use_persist_buf_on_alloc_fail=fsdp_db_use_persist_buf_on_alloc_fail,
         disable_symmetric_registration=disable_symmetric_registration,
+        megatron_fsdp_prefetch_recompute_forward_weights=prefetch_recompute_forward_weights,
+        megatron_fsdp_cache_param_bucket_views=cache_param_bucket_views,
         megatron_fsdp_use_decoupled_grad=use_decoupled_grad,
         megatron_fsdp_cuda_graph_mode=cuda_graph_mode,
     )
@@ -678,6 +689,8 @@ def fully_shard(
     fsdp_db_use_persist_buf_on_alloc_fail: bool = False,
     disable_symmetric_registration: bool = False,
     enable_fine_grained_param_gather: bool = False,
+    prefetch_recompute_forward_weights: bool = False,
+    cache_param_bucket_views: bool = False,
     use_decoupled_grad: bool = False,
     cuda_graph_mode: bool = False,
 ) -> tuple[MegatronFSDP, torch.optim.Optimizer]:
@@ -729,6 +742,9 @@ def fully_shard(
         fsdp_double_buffer=fsdp_double_buffer,
         fsdp_db_use_persist_buf_on_alloc_fail=fsdp_db_use_persist_buf_on_alloc_fail,
         disable_symmetric_registration=disable_symmetric_registration,
+        enable_fine_grained_param_gather=enable_fine_grained_param_gather,
+        prefetch_recompute_forward_weights=prefetch_recompute_forward_weights,
+        cache_param_bucket_views=cache_param_bucket_views,
         use_decoupled_grad=use_decoupled_grad,
         cuda_graph_mode=cuda_graph_mode,
     )

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

@@ -1,16 +1,4 @@
-# Copyright (c) 2025, NVIDIA CORPORATION.  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.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.import functools
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import functools
 import importlib
@@ -272,6 +260,9 @@ def __init__(
         self.enable_fine_grained_param_gather_backward_hook = (
             enable_fine_grained_param_gather_backward_hook
         )
+        self.prefetch_recompute_forward_weights = (
+            self.ddp_config.megatron_fsdp_prefetch_recompute_forward_weights
+        )
         self.report_nan_in_param_grad = report_nan_in_param_grad
 
         # FSDPDistributedIndex stores the process groups and meshes used by Megatron-FSDP.
@@ -879,9 +870,29 @@ def _pre_backward_param_unshard(module: nn.Module, *unused):
                 param_list = list(module.parameters(recurse=False))
 
             # All-gather / unshard the module parameters before the backward pass.
-            self.all_gather_and_wait_parameters_ready(
-                param_list, prefetch_order=PrefetchOrder.BACKWARD_PASS_ORDER, bwd=True
-            )
+            if self.prefetch_recompute_forward_weights:
+                self.all_gather_and_wait_parameters_ready(
+                    param_list,
+                    prefetch=False,
+                    prefetch_order=PrefetchOrder.BACKWARD_PASS_ORDER,
+                    bwd=True,
+                )
+                self.all_gather_and_wait_parameters_ready(
+                    param_list,
+                    prefetch=True,
+                    prefetch_order=PrefetchOrder.BACKWARD_PASS_ORDER,
+                    bwd=False,
+                )
+                self.all_gather_and_wait_parameters_ready(
+                    param_list,
+                    prefetch=True,
+                    prefetch_order=PrefetchOrder.BACKWARD_PASS_ORDER,
+                    bwd=True,
+                )
+            else:
+                self.all_gather_and_wait_parameters_ready(
+                    param_list, prefetch_order=PrefetchOrder.BACKWARD_PASS_ORDER, bwd=True
+                )
 
         self._root_pre_backward_hook_issued = False
 

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

@@ -1,16 +1,4 @@
-# Copyright (c) 2025, NVIDIA CORPORATION.  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.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 # TODO: Split this file into smaller files.
 
@@ -59,6 +47,18 @@
 logger = logging.getLogger(__name__)
 
 
+def _same_tensor_view(a: Optional[torch.Tensor], b: torch.Tensor) -> bool:
+    if a is None:
+        return False
+    return (
+        a.data_ptr() == b.data_ptr()
+        and a.dtype == b.dtype
+        and a.shape == b.shape
+        and a.stride() == b.stride()
+        and a.storage_offset() == b.storage_offset()
+    )
+
+
 try:
     # Default to Megatron-LM FW.
     from megatron.core.distributed.distributed_data_parallel_config import (
@@ -924,6 +924,8 @@ def __init__(
 
         # Count all parameters in this buffer and store their enumerated index.
         self.param_idx = {p: i for i, p in enumerate(self.params)}
+        self.cache_param_bucket_views = ddp_config.megatron_fsdp_cache_param_bucket_views
+        self._param_bucket_view_cache = {}
 
     def init_data(self, data: torch.Tensor):
         """Allocate a buffer Tensor to persistently store the data for this
@@ -970,6 +972,30 @@ def fetch_bucket(
 
         # Need to set parameter data after resize model weight buffer data-storage.
         if set_param_data:
+            self.set_param_data_from_bucket(bucket)
+        return bucket
+
+    def _bucket_view_cache_key(self, bucket: Bucket):
+        return (
+            bucket.data.data_ptr(),
+            bucket.data.numel(),
+            bucket.data.dtype,
+            str(bucket.data.device),
+            self.is_transpose_buffer,
+        )
+
+    def _build_param_bucket_view_entries(self, bucket: Bucket):
+        entries = []
+        for p in self.params:
+            item_id = self.param_idx[p]
+            p = to_local_if_dtensor(p)
+            data = self.get_item_from_bucket(bucket, item_id).view(p.shape)
+            entries.append((p, data, is_float8tensor(p)))
+        return entries
+
+    def set_param_data_from_bucket(self, bucket: Bucket) -> None:
+        """Attach module parameter tensors to their views in an all-gather bucket."""
+        if not self.cache_param_bucket_views:
             for p in self.params:
                 item_id = self.param_idx[p]
                 p = to_local_if_dtensor(p)
@@ -978,7 +1004,24 @@ def fetch_bucket(
                     fp8_set_raw_data(p, data, self.is_transpose_buffer)
                 else:
                     p.data = data
-        return bucket
+            return
+
+        cache_key = self._bucket_view_cache_key(bucket)
+        entries = self._param_bucket_view_cache.get(cache_key)
+        if entries is None:
+            entries = self._build_param_bucket_view_entries(bucket)
+            self._param_bucket_view_cache[cache_key] = entries
+
+        for p, data, is_fp8 in entries:
+            if is_fp8:
+                old_data = fp8_get_raw_data(p, self.is_transpose_buffer)
+                if _same_tensor_view(old_data, data):
+                    continue
+                fp8_set_raw_data(p, data, self.is_transpose_buffer)
+            else:
+                if _same_tensor_view(p.data, data):
+                    continue
+                p.data = data
 
     def allocate_bucket_storage(
         self,
@@ -4055,10 +4098,6 @@ def all_gather_params(
                     "but double buffers can support no more than 2 FSDP units."
                 )
 
-        # Do not release the buckets that are being all-gathered.
-        for bucket_id in ag_buckets:
-            self.bucket_can_be_released[self.get_bucket_key(bucket_id, bwd)] = False
-
         # If prefetch is enabled, we will add prefetch buckets to ag_buckets.
         if prefetch:
 
@@ -4128,6 +4167,11 @@ def need_skip_prefetch(bucket_id):
                 ag_buckets = list(sorted(set(ag_buckets)))
                 bucket_id = next_bucket_id(ag_buckets)
 
+        # Do not release the buckets that are requested by this call, even if
+        # they are already ready and do not need a new all-gather.
+        for bucket_id in ag_buckets:
+            self.bucket_can_be_released[self.get_bucket_key(bucket_id, bwd)] = False
+
         # Only all-gather on buckets that have not been allocated yet or whose
         # persistent storage was preserved but is not ready for use.
         ag_buckets = [