Skip to content

Preserve non-FSDP-unit buckets across AllGatherPipeline reset #4717

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wujingyue merged 11 commits into from
May 30, 2026
+267 −116
Merged

Preserve non-FSDP-unit buckets across AllGatherPipeline reset #4717

Show file tree
Hide file tree
Changes from 10 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
2dc9c5a
Preserve non-FSDP-unit buckets across AllGatherPipeline reset
wujingyue May 11, 2026
7f68d36
WIP
wujingyue May 11, 2026
2b4cf04
Update golden values
wujingyue May 12, 2026
1752fef
Add a repro
wujingyue May 12, 2026
766a028
Simplify
wujingyue May 13, 2026
3c81dbf
Fix lint
wujingyue May 13, 2026
e7c0cfc
Preserve non-FSDP buckets by default on reset
wujingyue May 19, 2026
7a46b4c
Rename preserved bucket status
wujingyue May 19, 2026
7c03222
Update FSDP overlap H100 golden memory
wujingyue May 21, 2026
6f9cf31
Clarify preserved non-unit bucket state
wujingyue May 28, 2026
c61795f
Merge branch 'main' into fsdp/preserve-non-unit-buckets-on-reset
wujingyue May 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1290,7 +1290,7 @@ def synchronize_param_gather(self):
"""
Synchronize parameter all-gather operations for all model parameters.
"""
self.all_gather_pipeline.reset()
self.all_gather_pipeline.reset(preserve_non_fsdp_units=True)
self._replace_param_with_distributed_if_needed()

def synchronize_gradient_reduce(self):
Expand Down
54 changes: 41 additions & 13 deletions megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3393,13 +3393,15 @@ class BucketStatus(Enum):

Attributes:
EMPTY (int): The bucket is empty and not in use.
PRESERVED (int): The bucket storage is retained but not ready for use.
COMMUNICATING (int): The bucket is currently being used for communication.
READY_TO_USE (int): The bucket is filled with data and ready for use.
"""

EMPTY = 1
COMMUNICATING = 2
READY_TO_USE = 3
PRESERVED = 2
COMMUNICATING = 3
READY_TO_USE = 4


class GradReducePipeline:
Expand Down Expand Up @@ -3921,8 +3923,14 @@ def num_buckets(self):
"""Return the number of buckets."""
return self.buffer.num_buckets

def reset(self):
"""Reset the pipeline state."""
def reset(self, preserve_non_fsdp_units: bool = True):
"""Reset the pipeline state.

Non-FSDP-unit buckets are preserved by default because their params may
be read across module boundaries. Setting preserve_non_fsdp_units=False
releases all bucket storage and is intended only for debugging when the
model will not be reused.
"""
if len(self.param_gather_event_map) > 0:
warnings.warn(
(
Expand All @@ -3934,12 +3942,26 @@ def reset(self):
while len(self.param_gather_event_map) > 0:
(bucket_id, bwd) = next(iter(self.param_gather_event_map))
self.wait_bucket_ready(bucket_id, bwd)

for bucket_id in range(self.num_buckets):
is_unit_bucket = self.buffer.parameter_groups[bucket_id].fsdp_unit_id is not None
for bwd in [False, True]:
self.bucket_can_be_released[self.get_bucket_key(bucket_id, bwd)] = True
bucket_key = self.get_bucket_key(bucket_id, bwd)
# If preserve_non_fsdp_units is set, then do not release buckets
# associated with FSDP non-units. Instead, mark the bucket as PRESERVED
# (not NEW) so a later all-gather refreshes preserved non-unit bucket
# storage in place.
if preserve_non_fsdp_units and not is_unit_bucket:
self.bucket_status[bucket_key] = BucketStatus.PRESERVED
else:
self.bucket_can_be_released[bucket_key] = True
self.recycle_unused_buckets()

assert all([status is BucketStatus.EMPTY for status in self.bucket_status.values()]), (
expected_statuses = (BucketStatus.EMPTY,)
if preserve_non_fsdp_units:
expected_statuses += (BucketStatus.PRESERVED,)

assert all(status in expected_statuses for status in self.bucket_status.values()), (
f"There are still working buckets, it is not safe to reset. "
f"bucket_status: {self.bucket_status}."
)
Expand Down Expand Up @@ -4067,11 +4089,13 @@ def need_skip_prefetch(bucket_id):
ag_buckets = list(sorted(set(ag_buckets)))
bucket_id = next_bucket_id(ag_buckets)

# Only all-gather on buckets that have not been allocated yet.
# 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 = [
bucket_id
for bucket_id in ag_buckets
if self.bucket_status[self.get_bucket_key(bucket_id, bwd)] == BucketStatus.EMPTY
if self.bucket_status[self.get_bucket_key(bucket_id, bwd)]
in (BucketStatus.EMPTY, BucketStatus.PRESERVED)
]
if len(ag_buckets) == 0:
return
Expand Down Expand Up @@ -4145,12 +4169,13 @@ def wait_bucket_ready(self, bucket_id, bwd, empty_ok=False):
if self.bucket_status[bucket_key] == BucketStatus.READY_TO_USE:
# Already ready to use.
return
if self.bucket_status[bucket_key] == BucketStatus.EMPTY:
if self.bucket_status[bucket_key] in (BucketStatus.EMPTY, BucketStatus.PRESERVED):
if empty_ok:
return
# Bucket shouldn't be empty, this implies that the bucket
# was not allocated or NCCL operations are not complete.
raise ValueError(f"Bucket {bucket_id} is empty.")
# Bucket should not be empty or merely preserved here; this implies that
# the bucket was not allocated, was not made ready for use, or NCCL
# operations are not complete.
raise ValueError(f"Bucket {bucket_id} is {self.bucket_status[bucket_key].name}.")

# Wait for asynchronous / overlapped NCCL operations to complete.
param_gather_event, mark_bucket_ready_to_use = self.param_gather_event_map.pop(bucket_key)
Expand Down Expand Up @@ -4241,7 +4266,10 @@ def async_bucket_gather(self, bucket_id, bwd) -> None:
bucket_key = self.get_bucket_key(bucket_id, bwd)

self.bucket_can_be_released[bucket_key] = False
if self.bucket_status[bucket_key] != BucketStatus.EMPTY:
if self.bucket_status[bucket_key] in (
BucketStatus.COMMUNICATING,
BucketStatus.READY_TO_USE,
):
return

self.bucket_status[bucket_key] = BucketStatus.COMMUNICATING
Expand Down
102 changes: 51 additions & 51 deletions ...tests/test_cases/moe/deepseek_proxy_fsdp_ep2_fsdp2_1node/golden_values_dev_dgx_gb200.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,56 +118,56 @@
"end_step": 50,
"step_interval": 1,
"values": {
"1": 594416640.0,
"2": 594418176.0,
"3": 594418176.0,
"4": 594418176.0,
"5": 594418176.0,
"6": 594418176.0,
"7": 594418176.0,
"8": 594418176.0,
"9": 594418176.0,
"10": 594418176.0,
"11": 594418176.0,
"12": 594418176.0,
"13": 594418176.0,
"14": 594418176.0,
"15": 594418176.0,
"16": 594418176.0,
"17": 594418176.0,
"18": 594418176.0,
"19": 594418176.0,
"20": 594418176.0,
"21": 594418176.0,
"22": 594418176.0,
"23": 594418176.0,
"24": 594418176.0,
"25": 594418176.0,
"26": 594418176.0,
"27": 594418176.0,
"28": 594418176.0,
"29": 594418176.0,
"30": 594418176.0,
"31": 594418176.0,
"32": 594418176.0,
"33": 594418176.0,
"34": 594418176.0,
"35": 594418176.0,
"36": 594418176.0,
"37": 594418176.0,
"38": 594418176.0,
"39": 594418176.0,
"40": 594418176.0,
"41": 594418176.0,
"42": 594418176.0,
"43": 594418176.0,
"44": 594418176.0,
"45": 594418176.0,
"46": 594418176.0,
"47": 594418176.0,
"48": 594418176.0,
"49": 594418176.0,
"50": 594418176.0
"1": 800470016.0,
"2": 800471552.0,
"3": 800471552.0,
"4": 800471552.0,
"5": 800471552.0,
"6": 800471552.0,
"7": 800471552.0,
"8": 800471552.0,
"9": 800471552.0,
"10": 800471552.0,
"11": 800471552.0,
"12": 800471552.0,
"13": 800471552.0,
"14": 800471552.0,
"15": 800471552.0,
"16": 800471552.0,
"17": 800471552.0,
"18": 800471552.0,
"19": 800471552.0,
"20": 800471552.0,
"21": 800471552.0,
"22": 800471552.0,
"23": 800471552.0,
"24": 800471552.0,
"25": 800471552.0,
"26": 800471552.0,
"27": 800471552.0,
"28": 800471552.0,
"29": 800471552.0,
"30": 800471552.0,
"31": 800471552.0,
"32": 800471552.0,
"33": 800471552.0,
"34": 800471552.0,
"35": 800471552.0,
"36": 800471552.0,
"37": 800471552.0,
"38": 800471552.0,
"39": 800471552.0,
"40": 800471552.0,
"41": 800471552.0,
"42": 800471552.0,
"43": 800471552.0,
"44": 800471552.0,
"45": 800471552.0,
"46": 800471552.0,
"47": 800471552.0,
"48": 800471552.0,
"49": 800471552.0,
"50": 800471552.0
}
},
"mem-max-allocated-bytes": {
Expand Down Expand Up @@ -284,4 +284,4 @@
"50": 1.09022
}
}
}
}
102 changes: 51 additions & 51 deletions ...s/test_cases/moe/deepseek_proxy_fsdp_ep2_fsdp2_ep_overlap/golden_values_dev_dgx_h100.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,56 +118,56 @@
"end_step": 50,
"step_interval": 1,
"values": {
"1": 407463424.0,
"2": 407515648.0,
"3": 407566336.0,
"4": 407617024.0,
"5": 407667712.0,
"6": 407718400.0,
"7": 407769088.0,
"8": 407819776.0,
"9": 407870464.0,
"10": 407921152.0,
"11": 407971840.0,
"12": 408022528.0,
"13": 408073216.0,
"14": 408123904.0,
"15": 408174592.0,
"16": 408225280.0,
"17": 408275968.0,
"18": 408326656.0,
"19": 408377344.0,
"20": 408428032.0,
"21": 408478720.0,
"22": 408529408.0,
"23": 408580096.0,
"24": 408630784.0,
"25": 408681472.0,
"26": 408732160.0,
"27": 408782848.0,
"28": 408833536.0,
"29": 408884224.0,
"30": 408934912.0,
"31": 408985600.0,
"32": 409036288.0,
"33": 409086976.0,
"34": 409137664.0,
"35": 409188352.0,
"36": 409239040.0,
"37": 409289728.0,
"38": 409340416.0,
"39": 409391104.0,
"40": 409441792.0,
"41": 409492480.0,
"42": 409543168.0,
"43": 409593856.0,
"44": 409644544.0,
"45": 409695232.0,
"46": 409745920.0,
"47": 409796608.0,
"48": 409847296.0,
"49": 409897984.0,
"50": 409948672.0
"1": 510494208.0,
"2": 510546432.0,
"3": 510597120.0,
"4": 510647808.0,
"5": 510698496.0,
"6": 510749184.0,
"7": 510799872.0,
"8": 510850560.0,
"9": 510901248.0,
"10": 510951936.0,
"11": 511002624.0,
"12": 511053312.0,
"13": 511104000.0,
"14": 511154688.0,
"15": 511205376.0,
"16": 511256064.0,
"17": 511306752.0,
"18": 511357440.0,
"19": 511408128.0,
"20": 511458816.0,
"21": 511509504.0,
"22": 511560192.0,
"23": 511610880.0,
"24": 511661568.0,
"25": 511712256.0,
"26": 511762944.0,
"27": 511813632.0,
"28": 511864320.0,
"29": 511915008.0,
"30": 511965696.0,
"31": 512016384.0,
"32": 512067072.0,
"33": 512117760.0,
"34": 512168448.0,
"35": 512219136.0,
"36": 512269824.0,
"37": 512320512.0,
"38": 512371200.0,
"39": 512421888.0,
"40": 512472576.0,
"41": 512523264.0,
"42": 512573952.0,
"43": 512624640.0,
"44": 512675328.0,
"45": 512726016.0,
"46": 512776704.0,
"47": 512827392.0,
"48": 512878080.0,
"49": 512928768.0,
"50": 512979456.0
}
},
"mem-max-allocated-bytes": {
Expand Down Expand Up @@ -284,4 +284,4 @@
"50": 0.05891
}
}
}
}
Loading
Loading