From 4c1013a45b464cb22d267c3c5ed0cec013b0170a Mon Sep 17 00:00:00 2001 From: hongbinl Date: Mon, 1 Jun 2026 06:49:03 -0700 Subject: [PATCH 01/12] Optimize FSDP param bucket view setup (cherry picked from commit bb3092d44b8abeb526a1b1833a1e340b09e02c6d) Signed-off-by: hongbinl --- .../megatron_fsdp/param_and_grad_buffer.py | 88 +++++++++++++++++-- 1 file changed, 79 insertions(+), 9 deletions(-) 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 2d7f573afdb..5eedda1d05c 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 @@ -21,6 +21,7 @@ import inspect import logging import math +import os import traceback import warnings from collections import defaultdict, namedtuple @@ -59,6 +60,22 @@ logger = logging.getLogger(__name__) +def _env_flag(name: str) -> bool: + return os.environ.get(name, "").lower() in ("1", "true", "yes", "on") + + +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 +941,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 = _env_flag("MCORE_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,16 +989,50 @@ def fetch_bucket( # Need to set parameter data after resize model weight buffer data-storage. if set_param_data: - 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) - if is_float8tensor(p): - fp8_set_raw_data(p, data, self.is_transpose_buffer) - else: - p.data = 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.""" + entries = None + if self.cache_param_bucket_views: + 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 + else: + entries = self._build_param_bucket_view_entries(bucket) + + for p, data, is_fp8 in entries: + if is_fp8: + old_data = fp8_get_raw_data(p, self.is_transpose_buffer) + if self.cache_param_bucket_views and _same_tensor_view(old_data, data): + continue + fp8_set_raw_data(p, data, self.is_transpose_buffer) + else: + if self.cache_param_bucket_views and _same_tensor_view(p.data, data): + continue + p.data = data + def allocate_bucket_storage( self, shard: bool = False, @@ -3911,6 +3964,8 @@ def __init__( for i in range(self.buffer.num_buckets): for bwd in [False, True]: self.bucket_can_be_released[self.get_bucket_key(i, bwd)] = False + self.defer_param_bucket_view_setup = _env_flag("MCORE_FSDP_DEFER_PARAM_VIEW_SETUP") + self.deferred_param_bucket_views = {} # Map each bucket to the bucket group it belongs to by enumerated ID. # Made to collect a subset of buckets in the same bucket group. @@ -4150,6 +4205,10 @@ def need_skip_prefetch(bucket_id): # into an allocated bucket containing unsharded weights. self.async_bucket_gather(bucket_id, bwd) + if self.defer_param_bucket_view_setup: + for bucket_id in buckets: + self.set_deferred_param_bucket_views(bucket_id, bwd) + # Replace the parameter all-gather event with coalescing event. for bucket_id in buckets: bucket_key = self.get_bucket_key(bucket_id, bwd) @@ -4245,6 +4304,15 @@ def recycle_unused_buckets(self): self.release_bucket(bucket_id, is_transpose_weight) self.bucket_can_be_released[bucket_key] = False + def set_deferred_param_bucket_views(self, bucket_id: int, bwd: bool) -> None: + """Attach parameter views after the all-gather has been enqueued.""" + bucket_key = self.get_bucket_key(bucket_id, bwd) + pending = self.deferred_param_bucket_views.pop(bucket_key, None) + if pending is None: + return + wbuf, bucket = pending + wbuf.set_param_data_from_bucket(bucket) + def get_fsdp_buffer(self, bucket_id: int, bwd=False) -> DataParallelBuffer: """ Get the FSDP / DP-Shard buffer with the given bucket ID. @@ -4280,7 +4348,9 @@ def async_bucket_gather(self, bucket_id, bwd) -> None: self.recycle_unused_buckets() # Allocate an empty bucket to store the module weights. - bucket = wbuf.fetch_bucket(set_param_data=True) + bucket = wbuf.fetch_bucket(set_param_data=not self.defer_param_bucket_view_setup) + if self.defer_param_bucket_view_setup: + self.deferred_param_bucket_views[bucket_key] = (wbuf, bucket) # All-gather the module weights in each buffer shard into the allocated bucket. # Now each rank will have a copy of this FSDP unit module's weights. From 0d4016e7e059c5b0ed1db612fc8100c66bb76b3d Mon Sep 17 00:00:00 2001 From: hongbinl Date: Mon, 1 Jun 2026 19:56:45 -0700 Subject: [PATCH 02/12] Prefetch recompute forward FSDP weights (cherry picked from commit f0c0aeeaad7ea783e049b549e5587c2d131456fd) Conflicts: megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py Signed-off-by: hongbinl --- .../fsdp/src/megatron_fsdp/megatron_fsdp.py | 46 +++++++++++++++++-- .../megatron_fsdp/param_and_grad_buffer.py | 27 ++++++++--- 2 files changed, 63 insertions(+), 10 deletions(-) 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 1c71cd33c74..f1dd6609343 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -15,6 +15,7 @@ import functools import importlib import logging +import os from contextlib import contextmanager from enum import Enum, auto from functools import partial @@ -272,6 +273,9 @@ def __init__( self.enable_fine_grained_param_gather_backward_hook = ( enable_fine_grained_param_gather_backward_hook ) + self.prefetch_recompute_forward_weights = os.environ.get( + "MCORE_FSDP_PREFETCH_RECOMPUTE_FORWARD_WEIGHTS", "0" + ).lower() in ("1", "true", "yes", "on") self.report_nan_in_param_grad = report_nan_in_param_grad # FSDPDistributedIndex stores the process groups and meshes used by Megatron-FSDP. @@ -878,9 +882,45 @@ 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, + ) + + outer_fsdp_group_param_gather = ( + self.dist_index.use_hybrid_fsdp + and self.ddp_config.outer_dp_sharding_strategy != "no_shard" + and (self.microbatch_count == 0 or self.model_auto_sync) + ) + # During full activation recomputation, the next backward layer + # first reruns its forward path, so row-wise weights are more + # urgent than the column-wise transpose buffers used later in + # the same layer's backward path. + self.all_gather_pipeline.all_gather_params( + params=param_list, + prefetch=True, + prefetch_order=PrefetchOrder.BACKWARD_PASS_ORDER, + suggested_AG_prefetch_size=self.suggested_AG_prefetch_size, + outer_fsdp_group_param_gather=outer_fsdp_group_param_gather, + bwd=False, + include_current=False, + ) + self.all_gather_pipeline.all_gather_params( + params=param_list, + prefetch=True, + prefetch_order=PrefetchOrder.BACKWARD_PASS_ORDER, + suggested_AG_prefetch_size=self.suggested_AG_prefetch_size, + outer_fsdp_group_param_gather=outer_fsdp_group_param_gather, + bwd=True, + include_current=False, + ) + 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 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 5eedda1d05c..e4675ba9883 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 @@ -4040,6 +4040,7 @@ def all_gather_params( async_param_gather: bool = True, outer_fsdp_group_param_gather: bool = False, bwd: bool = False, + include_current: bool = True, ): """All-gather the params. If prefetch is enabled, prefetch next buckets in the order of `prefetch_order`. @@ -4056,16 +4057,22 @@ def all_gather_params( bwd (bool, optional): Whether to all-gather column-wise parameters instead of row-wise parameters for the backward pass for formats that require a transpose buffer like MXFP8. + include_current (bool, optional): + Whether to all-gather the buckets that contain ``params``. If False, + only the prefetch buckets adjacent to the current buckets are issued. """ if len(params) == 0: return - ag_buckets = [self.buffer.param_to_param_group[item] for item in params] - ag_buckets = list(sorted(set(ag_buckets))) # Sort in order of unique bucket ID. + current_ag_buckets = [self.buffer.param_to_param_group[item] for item in params] + current_ag_buckets = list( + sorted(set(current_ag_buckets)) + ) # Sort in order of unique bucket ID. + ag_buckets = list(current_ag_buckets) parameter_groups = self.buffer.parameter_groups if self.buffer.ddp_config.fsdp_double_buffer: double_buf_units = set() - for bucket_id in ag_buckets: + for bucket_id in current_ag_buckets: fsdp_unit_id = parameter_groups[bucket_id].fsdp_unit_id if fsdp_unit_id in self.buffer.double_buf_units: double_buf_units.add(fsdp_unit_id) @@ -4075,10 +4082,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: @@ -4148,6 +4151,12 @@ def need_skip_prefetch(bucket_id): ag_buckets = list(sorted(set(ag_buckets))) bucket_id = next_bucket_id(ag_buckets) + if not include_current: + current_ag_bucket_set = set(current_ag_buckets) + ag_buckets = [ + bucket_id for bucket_id in ag_buckets if bucket_id not in current_ag_bucket_set + ] + # Only all-gather on buckets that have not been allocated yet. ag_buckets = [ bucket_id @@ -4157,6 +4166,10 @@ def need_skip_prefetch(bucket_id): if len(ag_buckets) == 0: return + # 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 + # Divide buckets into aggregate groups. We need to reconstruct the bucket groups # because the all-gather parameter groups may be a subset of the buckets. bucket_group_to_buckets = {} From 079a5a346e0fe981e0ae0db9a081ef6b8929cebb Mon Sep 17 00:00:00 2001 From: hongbinl Date: Thu, 4 Jun 2026 06:38:07 -0700 Subject: [PATCH 03/12] chore: update FSDP copyright headers Signed-off-by: hongbinl --- .../fsdp/src/megatron_fsdp/megatron_fsdp.py | 14 +------------- .../src/megatron_fsdp/param_and_grad_buffer.py | 14 +------------- 2 files changed, 2 insertions(+), 26 deletions(-) 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 f1dd6609343..7bbaf99e91e 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/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 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 e4675ba9883..42db91a2eb4 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 @@ -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. From bef8c71b0f361ac2dd8738d1771655868dc7a28d Mon Sep 17 00:00:00 2001 From: hongbinl Date: Thu, 4 Jun 2026 07:10:56 -0700 Subject: [PATCH 04/12] Fix FSDP all-gather bucket release protection Conflicts: megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py Signed-off-by: hongbinl --- .../fsdp/src/megatron_fsdp/param_and_grad_buffer.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) 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 42db91a2eb4..5292d62e01d 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 @@ -4145,6 +4145,11 @@ def need_skip_prefetch(bucket_id): bucket_id for bucket_id in ag_buckets if bucket_id not in current_ag_bucket_set ] + # 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. ag_buckets = [ bucket_id @@ -4154,10 +4159,6 @@ def need_skip_prefetch(bucket_id): if len(ag_buckets) == 0: return - # 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 - # Divide buckets into aggregate groups. We need to reconstruct the bucket groups # because the all-gather parameter groups may be a subset of the buckets. bucket_group_to_buckets = {} From fb698f43efc264bbc1cb4a05be99f403034fadff Mon Sep 17 00:00:00 2001 From: hongbinl Date: Thu, 4 Jun 2026 07:27:44 -0700 Subject: [PATCH 05/12] Use config flag for recompute weight prefetch Conflicts: megatron/core/distributed/distributed_data_parallel_config.py megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py megatron/training/arguments.py Signed-off-by: hongbinl --- .../distributed_data_parallel_config.py | 6 ++++ .../distributed_data_parallel_config.py | 6 ++++ .../fsdp/src/megatron_fsdp/fully_shard.py | 9 ++++++ .../fsdp/src/megatron_fsdp/megatron_fsdp.py | 7 ++--- .../megatron_fsdp/param_and_grad_buffer.py | 28 +++---------------- megatron/training/arguments.py | 12 ++++++++ 6 files changed, 40 insertions(+), 28 deletions(-) diff --git a/megatron/core/distributed/distributed_data_parallel_config.py b/megatron/core/distributed/distributed_data_parallel_config.py index 50878e149de..d5dd26a9868 100644 --- a/megatron/core/distributed/distributed_data_parallel_config.py +++ b/megatron/core/distributed/distributed_data_parallel_config.py @@ -206,6 +206,12 @@ 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. This + also caches parameter bucket views to reduce repeated Python-side view setup. + """ + def __post_init__(self): import os diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py index a2feb99cb23..183fb192a80 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py @@ -151,6 +151,12 @@ 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. This + also caches parameter bucket views to reduce repeated Python-side view setup. + """ + def __post_init__(self): import os diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py index a1f7fabd50a..e3ded22f94d 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py @@ -103,6 +103,7 @@ 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, use_decoupled_grad: bool = False, ) -> torch.nn.Module: """ @@ -261,6 +262,11 @@ 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. This also caches + parameter bucket views to reduce repeated Python-side view setup. Defaults to False. + use_decoupled_grad (bool): If true, reduced gradients are installed into `Parameter.decoupled_grad` instead of `Parameter.grad`. Defaults to False. @@ -359,6 +365,7 @@ 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_use_decoupled_grad=use_decoupled_grad, ) @@ -665,6 +672,7 @@ 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, use_decoupled_grad: bool = False, ) -> tuple[MegatronFSDP, torch.optim.Optimizer]: """ @@ -716,6 +724,7 @@ def fully_shard( 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, use_decoupled_grad=use_decoupled_grad, ) 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 7bbaf99e91e..99b82bc7d42 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -3,7 +3,6 @@ import functools import importlib import logging -import os from contextlib import contextmanager from enum import Enum, auto from functools import partial @@ -261,9 +260,9 @@ def __init__( self.enable_fine_grained_param_gather_backward_hook = ( enable_fine_grained_param_gather_backward_hook ) - self.prefetch_recompute_forward_weights = os.environ.get( - "MCORE_FSDP_PREFETCH_RECOMPUTE_FORWARD_WEIGHTS", "0" - ).lower() in ("1", "true", "yes", "on") + 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. 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 5292d62e01d..1dffc9ac74f 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 @@ -9,7 +9,6 @@ import inspect import logging import math -import os import traceback import warnings from collections import defaultdict, namedtuple @@ -48,10 +47,6 @@ logger = logging.getLogger(__name__) -def _env_flag(name: str) -> bool: - return os.environ.get(name, "").lower() in ("1", "true", "yes", "on") - - def _same_tensor_view(a: Optional[torch.Tensor], b: torch.Tensor) -> bool: if a is None: return False @@ -929,7 +924,9 @@ 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 = _env_flag("MCORE_FSDP_CACHE_PARAM_BUCKET_VIEWS") + self.cache_param_bucket_views = ( + ddp_config.megatron_fsdp_prefetch_recompute_forward_weights + ) self._param_bucket_view_cache = {} def init_data(self, data: torch.Tensor): @@ -3952,8 +3949,6 @@ def __init__( for i in range(self.buffer.num_buckets): for bwd in [False, True]: self.bucket_can_be_released[self.get_bucket_key(i, bwd)] = False - self.defer_param_bucket_view_setup = _env_flag("MCORE_FSDP_DEFER_PARAM_VIEW_SETUP") - self.deferred_param_bucket_views = {} # Map each bucket to the bucket group it belongs to by enumerated ID. # Made to collect a subset of buckets in the same bucket group. @@ -4207,10 +4202,6 @@ def need_skip_prefetch(bucket_id): # into an allocated bucket containing unsharded weights. self.async_bucket_gather(bucket_id, bwd) - if self.defer_param_bucket_view_setup: - for bucket_id in buckets: - self.set_deferred_param_bucket_views(bucket_id, bwd) - # Replace the parameter all-gather event with coalescing event. for bucket_id in buckets: bucket_key = self.get_bucket_key(bucket_id, bwd) @@ -4306,15 +4297,6 @@ def recycle_unused_buckets(self): self.release_bucket(bucket_id, is_transpose_weight) self.bucket_can_be_released[bucket_key] = False - def set_deferred_param_bucket_views(self, bucket_id: int, bwd: bool) -> None: - """Attach parameter views after the all-gather has been enqueued.""" - bucket_key = self.get_bucket_key(bucket_id, bwd) - pending = self.deferred_param_bucket_views.pop(bucket_key, None) - if pending is None: - return - wbuf, bucket = pending - wbuf.set_param_data_from_bucket(bucket) - def get_fsdp_buffer(self, bucket_id: int, bwd=False) -> DataParallelBuffer: """ Get the FSDP / DP-Shard buffer with the given bucket ID. @@ -4350,9 +4332,7 @@ def async_bucket_gather(self, bucket_id, bwd) -> None: self.recycle_unused_buckets() # Allocate an empty bucket to store the module weights. - bucket = wbuf.fetch_bucket(set_param_data=not self.defer_param_bucket_view_setup) - if self.defer_param_bucket_view_setup: - self.deferred_param_bucket_views[bucket_key] = (wbuf, bucket) + bucket = wbuf.fetch_bucket(set_param_data=True) # All-gather the module weights in each buffer shard into the allocated bucket. # Now each rank will have a copy of this FSDP unit module's weights. diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 423a51f061d..42a73ecf1ff 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -4783,6 +4783,18 @@ def _add_experimental_args(parser): "be used for the gradient communication / reduction data-type. When using NCCL " "v2.27+, reduction is always computed in FP32 if using NCCL Symmetric kernels.", ) + group.add_argument( + '--megatron-fsdp-prefetch-recompute-forward-weights', + action='store_true', + default=False, + dest='megatron_fsdp_prefetch_recompute_forward_weights', + help=( + 'If set, Megatron-FSDP prefetches rowwise weights needed by activation ' + 'recomputation during backward before prefetching backward transpose ' + 'weights. This also caches parameter bucket views to reduce repeated ' + 'Python-side view setup.' + ), + ) return parser From 6122a4540c06916cda6cb2ea78bb49deadf33cfd Mon Sep 17 00:00:00 2001 From: hongbinl Date: Thu, 4 Jun 2026 07:59:10 -0700 Subject: [PATCH 06/12] Validate FSDP recompute prefetch flag Conflicts: megatron/core/distributed/distributed_data_parallel_config.py megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py megatron/training/arguments.py Signed-off-by: hongbinl --- .../distributed_data_parallel_config.py | 9 +++++++++ .../distributed_data_parallel_config.py | 6 ++++++ .../fsdp/src/megatron_fsdp/fully_shard.py | 11 +++++++++++ megatron/training/arguments.py | 19 +++++++++++++++++++ 4 files changed, 45 insertions(+) diff --git a/megatron/core/distributed/distributed_data_parallel_config.py b/megatron/core/distributed/distributed_data_parallel_config.py index d5dd26a9868..fd018833c13 100644 --- a/megatron/core/distributed/distributed_data_parallel_config.py +++ b/megatron/core/distributed/distributed_data_parallel_config.py @@ -219,6 +219,15 @@ 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.nccl_ub: if 'expandable_segments:True' in os.getenv('PYTORCH_CUDA_ALLOC_CONF', '').split(','): raise ValueError( diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py index 183fb192a80..10bb370ff50 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py @@ -161,6 +161,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: if 'expandable_segments:True' in os.getenv('PYTORCH_CUDA_ALLOC_CONF', '').split(','): raise ValueError( diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py index e3ded22f94d..6bd0b1612a1 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py @@ -352,6 +352,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( diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 42a73ecf1ff..865c38711d0 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1268,6 +1268,25 @@ def validate_args(args, defaults={}): args.ckpt_format == "fsdp_dtensor" ), "Megatron-FSDP requires the `fsdp_dtensor` checkpointing format." + if args.megatron_fsdp_prefetch_recompute_forward_weights: + assert args.data_parallel_sharding_strategy == "optim_grads_params", ( + "--megatron-fsdp-prefetch-recompute-forward-weights is only supported " + 'with --data-parallel-sharding-strategy optim_grads_params.' + ) + assert args.recompute_granularity == "full", ( + "--megatron-fsdp-prefetch-recompute-forward-weights is only supported " + "with full activation recomputation." + ) + assert not args.overlap_moe_expert_parallel_comm, ( + "--megatron-fsdp-prefetch-recompute-forward-weights is not supported " + "with --overlap-moe-expert-parallel-comm." + ) + else: + assert not args.megatron_fsdp_prefetch_recompute_forward_weights, ( + "--megatron-fsdp-prefetch-recompute-forward-weights requires " + "--use-megatron-fsdp." + ) + if args.nccl_ub and args.use_megatron_fsdp: # In Megatron-LM, required implementation for manual registration is already provided. # So we enable the manual registration by default when nccl-ub and use_megatron_fsdp is set. From 3745d139e875123b230f04b3c4b91a6514ee1e5d Mon Sep 17 00:00:00 2001 From: hongbinl Date: Thu, 4 Jun 2026 19:10:11 -0700 Subject: [PATCH 07/12] chore: update FSDP dev copyright headers Signed-off-by: hongbinl --- .../distributed_data_parallel_config.py | 2 +- .../distributed_data_parallel_config.py | 2 +- .../fsdp/src/megatron_fsdp/fully_shard.py | 14 +------------- 3 files changed, 3 insertions(+), 15 deletions(-) diff --git a/megatron/core/distributed/distributed_data_parallel_config.py b/megatron/core/distributed/distributed_data_parallel_config.py index fd018833c13..55d2886c5a7 100644 --- a/megatron/core/distributed/distributed_data_parallel_config.py +++ b/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 diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py index 10bb370ff50..5aa6cce90c7 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py +++ b/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 diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py index 6bd0b1612a1..3794cff4f91 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py +++ b/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 From bcf9e8caa89c5735a4c7886b5e032ea171e1310e Mon Sep 17 00:00:00 2001 From: hongbinl Date: Thu, 4 Jun 2026 21:46:07 -0700 Subject: [PATCH 08/12] chore: fix FSDP recompute prefetch lint Signed-off-by: hongbinl --- .../core/distributed/distributed_data_parallel_config.py | 6 +++--- .../fsdp/src/megatron_fsdp/param_and_grad_buffer.py | 4 +--- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/megatron/core/distributed/distributed_data_parallel_config.py b/megatron/core/distributed/distributed_data_parallel_config.py index 55d2886c5a7..f0313f7f948 100644 --- a/megatron/core/distributed/distributed_data_parallel_config.py +++ b/megatron/core/distributed/distributed_data_parallel_config.py @@ -220,9 +220,9 @@ def __post_init__(self): 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.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'." 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 1dffc9ac74f..255bebd8522 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 @@ -924,9 +924,7 @@ 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_prefetch_recompute_forward_weights - ) + self.cache_param_bucket_views = ddp_config.megatron_fsdp_prefetch_recompute_forward_weights self._param_bucket_view_cache = {} def init_data(self, data: torch.Tensor): From c9b9e63d43b100c4f5cdb8440643b7748faebedf Mon Sep 17 00:00:00 2001 From: hongbinl Date: Mon, 8 Jun 2026 04:06:16 -0700 Subject: [PATCH 09/12] refactor: simplify FSDP recompute prefetch scheduling Signed-off-by: hongbinl --- .../fsdp/src/megatron_fsdp/megatron_fsdp.py | 24 ++++--------------- 1 file changed, 4 insertions(+), 20 deletions(-) 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 99b82bc7d42..848508a4359 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -876,33 +876,17 @@ def _pre_backward_param_unshard(module: nn.Module, *unused): prefetch_order=PrefetchOrder.BACKWARD_PASS_ORDER, bwd=True, ) - - outer_fsdp_group_param_gather = ( - self.dist_index.use_hybrid_fsdp - and self.ddp_config.outer_dp_sharding_strategy != "no_shard" - and (self.microbatch_count == 0 or self.model_auto_sync) - ) - # During full activation recomputation, the next backward layer - # first reruns its forward path, so row-wise weights are more - # urgent than the column-wise transpose buffers used later in - # the same layer's backward path. - self.all_gather_pipeline.all_gather_params( - params=param_list, + self.all_gather_and_wait_parameters_ready( + param_list, prefetch=True, prefetch_order=PrefetchOrder.BACKWARD_PASS_ORDER, - suggested_AG_prefetch_size=self.suggested_AG_prefetch_size, - outer_fsdp_group_param_gather=outer_fsdp_group_param_gather, bwd=False, - include_current=False, ) - self.all_gather_pipeline.all_gather_params( - params=param_list, + self.all_gather_and_wait_parameters_ready( + param_list, prefetch=True, prefetch_order=PrefetchOrder.BACKWARD_PASS_ORDER, - suggested_AG_prefetch_size=self.suggested_AG_prefetch_size, - outer_fsdp_group_param_gather=outer_fsdp_group_param_gather, bwd=True, - include_current=False, ) else: self.all_gather_and_wait_parameters_ready( From 8149d917775837a50f9702f567bddfc2b7380aa5 Mon Sep 17 00:00:00 2001 From: hongbinl Date: Mon, 8 Jun 2026 06:06:36 -0700 Subject: [PATCH 10/12] refactor: remove unused FSDP prefetch include-current option Signed-off-by: hongbinl --- .../fsdp/src/megatron_fsdp/param_and_grad_buffer.py | 10 ---------- 1 file changed, 10 deletions(-) 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 255bebd8522..a4a4808f678 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 @@ -4021,7 +4021,6 @@ def all_gather_params( async_param_gather: bool = True, outer_fsdp_group_param_gather: bool = False, bwd: bool = False, - include_current: bool = True, ): """All-gather the params. If prefetch is enabled, prefetch next buckets in the order of `prefetch_order`. @@ -4038,9 +4037,6 @@ def all_gather_params( bwd (bool, optional): Whether to all-gather column-wise parameters instead of row-wise parameters for the backward pass for formats that require a transpose buffer like MXFP8. - include_current (bool, optional): - Whether to all-gather the buckets that contain ``params``. If False, - only the prefetch buckets adjacent to the current buckets are issued. """ if len(params) == 0: return @@ -4132,12 +4128,6 @@ def need_skip_prefetch(bucket_id): ag_buckets = list(sorted(set(ag_buckets))) bucket_id = next_bucket_id(ag_buckets) - if not include_current: - current_ag_bucket_set = set(current_ag_buckets) - ag_buckets = [ - bucket_id for bucket_id in ag_buckets if bucket_id not in current_ag_bucket_set - ] - # 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: From 94355a4b5e20f0f5f857a2e76817992748fe0693 Mon Sep 17 00:00:00 2001 From: hongbinl Date: Mon, 8 Jun 2026 06:22:19 -0700 Subject: [PATCH 11/12] refactor: remove redundant FSDP current bucket list Signed-off-by: hongbinl --- .../fsdp/src/megatron_fsdp/param_and_grad_buffer.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) 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 a4a4808f678..40e4341fad4 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 @@ -4041,15 +4041,12 @@ def all_gather_params( if len(params) == 0: return - current_ag_buckets = [self.buffer.param_to_param_group[item] for item in params] - current_ag_buckets = list( - sorted(set(current_ag_buckets)) - ) # Sort in order of unique bucket ID. - ag_buckets = list(current_ag_buckets) + ag_buckets = [self.buffer.param_to_param_group[item] for item in params] + ag_buckets = list(sorted(set(ag_buckets))) # Sort in order of unique bucket ID. parameter_groups = self.buffer.parameter_groups if self.buffer.ddp_config.fsdp_double_buffer: double_buf_units = set() - for bucket_id in current_ag_buckets: + for bucket_id in ag_buckets: fsdp_unit_id = parameter_groups[bucket_id].fsdp_unit_id if fsdp_unit_id in self.buffer.double_buf_units: double_buf_units.add(fsdp_unit_id) From 26c95cde34e99393772d3b0c9c8f801af7f0d062 Mon Sep 17 00:00:00 2001 From: hongbinl Date: Fri, 12 Jun 2026 06:52:12 -0700 Subject: [PATCH 12/12] refactor: split FSDP bucket view cache flag Signed-off-by: hongbinl --- .../distributed_data_parallel_config.py | 13 +++++++-- .../distributed_data_parallel_config.py | 8 +++-- .../fsdp/src/megatron_fsdp/fully_shard.py | 11 +++++-- .../megatron_fsdp/param_and_grad_buffer.py | 29 ++++++++++++------- megatron/training/arguments.py | 16 ++++++++-- 5 files changed, 58 insertions(+), 19 deletions(-) diff --git a/megatron/core/distributed/distributed_data_parallel_config.py b/megatron/core/distributed/distributed_data_parallel_config.py index 95e4a6fba2e..ed2c43a2d4f 100644 --- a/megatron/core/distributed/distributed_data_parallel_config.py +++ b/megatron/core/distributed/distributed_data_parallel_config.py @@ -217,8 +217,12 @@ class DistributedDataParallelConfig: 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. This - also caches parameter bucket views to reduce repeated Python-side view setup. + 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 @@ -259,6 +263,11 @@ def __post_init__(self): "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( diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py index f80774abd67..938e17a5b3f 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py @@ -157,8 +157,12 @@ class DistributedDataParallelConfig: 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. This - also caches parameter bucket views to reduce repeated Python-side view setup. + 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 diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py index e9338ee55e5..fc6367bc02b 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py @@ -92,6 +92,7 @@ def fully_shard_model( 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: @@ -253,8 +254,11 @@ class that schedules the sharding lifecycle of the model parameters and gradient prefetch_recompute_forward_weights (bool): Whether to prefetch rowwise weights needed by activation recomputation during - backward before prefetching backward transpose weights. This also caches - parameter bucket views to reduce repeated Python-side view setup. Defaults to False. + 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 @@ -377,6 +381,7 @@ class that schedules the sharding lifecycle of the model parameters and gradient 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, ) @@ -685,6 +690,7 @@ def fully_shard( 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]: @@ -738,6 +744,7 @@ def fully_shard( 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, ) 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 8858e736e98..c81ca9e34c1 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 @@ -924,7 +924,7 @@ 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_prefetch_recompute_forward_weights + 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): @@ -995,24 +995,31 @@ def _build_param_bucket_view_entries(self, bucket: Bucket): def set_param_data_from_bucket(self, bucket: Bucket) -> None: """Attach module parameter tensors to their views in an all-gather bucket.""" - entries = None - if self.cache_param_bucket_views: - 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 - else: + if not self.cache_param_bucket_views: + 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) + if is_float8tensor(p): + fp8_set_raw_data(p, data, self.is_transpose_buffer) + else: + p.data = data + 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 self.cache_param_bucket_views and _same_tensor_view(old_data, data): + if _same_tensor_view(old_data, data): continue fp8_set_raw_data(p, data, self.is_transpose_buffer) else: - if self.cache_param_bucket_views and _same_tensor_view(p.data, data): + if _same_tensor_view(p.data, data): continue p.data = data diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 440cabcb17c..14ab59b71d3 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1297,6 +1297,9 @@ def validate_args(args, defaults={}): assert not args.megatron_fsdp_prefetch_recompute_forward_weights, ( "--megatron-fsdp-prefetch-recompute-forward-weights requires " "--use-megatron-fsdp." ) + assert not args.megatron_fsdp_cache_param_bucket_views, ( + "--megatron-fsdp-cache-param-bucket-views requires " "--use-megatron-fsdp." + ) if args.nccl_ub and args.use_megatron_fsdp: # In Megatron-LM, required implementation for manual registration is already provided. @@ -4906,8 +4909,17 @@ def _add_experimental_args(parser): help=( 'If set, Megatron-FSDP prefetches rowwise weights needed by activation ' 'recomputation during backward before prefetching backward transpose ' - 'weights. This also caches parameter bucket views to reduce repeated ' - 'Python-side view setup.' + 'weights.' + ), + ) + group.add_argument( + '--megatron-fsdp-cache-param-bucket-views', + action='store_true', + default=False, + dest='megatron_fsdp_cache_param_bucket_views', + help=( + 'If set, Megatron-FSDP caches parameter bucket views to reduce repeated ' + 'Python-side view setup when attaching module parameters to all-gather buckets.' ), ) group.add_argument(