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 72e04c0ba19..37777e7581b 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -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): 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 3e1b2bfc7fb..690ec263890 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 @@ -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: @@ -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( ( @@ -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}." ) @@ -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 @@ -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) @@ -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 diff --git a/tests/functional_tests/test_cases/moe/deepseek_proxy_fsdp_ep2_fsdp2_1node/golden_values_dev_dgx_gb200.json b/tests/functional_tests/test_cases/moe/deepseek_proxy_fsdp_ep2_fsdp2_1node/golden_values_dev_dgx_gb200.json index 34abdb88474..b3109ea4e60 100644 --- a/tests/functional_tests/test_cases/moe/deepseek_proxy_fsdp_ep2_fsdp2_1node/golden_values_dev_dgx_gb200.json +++ b/tests/functional_tests/test_cases/moe/deepseek_proxy_fsdp_ep2_fsdp2_1node/golden_values_dev_dgx_gb200.json @@ -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": { @@ -284,4 +284,4 @@ "50": 1.09022 } } -} \ No newline at end of file +} diff --git a/tests/functional_tests/test_cases/moe/deepseek_proxy_fsdp_ep2_fsdp2_ep_overlap/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/moe/deepseek_proxy_fsdp_ep2_fsdp2_ep_overlap/golden_values_dev_dgx_h100.json index b18bfdc4c40..3d20587a302 100644 --- a/tests/functional_tests/test_cases/moe/deepseek_proxy_fsdp_ep2_fsdp2_ep_overlap/golden_values_dev_dgx_h100.json +++ b/tests/functional_tests/test_cases/moe/deepseek_proxy_fsdp_ep2_fsdp2_ep_overlap/golden_values_dev_dgx_h100.json @@ -118,56 +118,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 411756032.0, - "2": 411808256.0, - "3": 411858944.0, - "4": 411909632.0, - "5": 411960320.0, - "6": 412011008.0, - "7": 412061696.0, - "8": 412112384.0, - "9": 412163072.0, - "10": 412213760.0, - "11": 412264448.0, - "12": 412315136.0, - "13": 412365824.0, - "14": 412416512.0, - "15": 412467200.0, - "16": 412517888.0, - "17": 412568576.0, - "18": 412619264.0, - "19": 412669952.0, - "20": 412720640.0, - "21": 412771328.0, - "22": 412822016.0, - "23": 412872704.0, - "24": 412923392.0, - "25": 412974080.0, - "26": 413024768.0, - "27": 413075456.0, - "28": 413126144.0, - "29": 413176832.0, - "30": 413227520.0, - "31": 413278208.0, - "32": 413328896.0, - "33": 413379584.0, - "34": 413430272.0, - "35": 413480960.0, - "36": 413531648.0, - "37": 413582336.0, - "38": 413633024.0, - "39": 413683712.0, - "40": 413734400.0, - "41": 413785088.0, - "42": 413835776.0, - "43": 413886464.0, - "44": 413937152.0, - "45": 413987840.0, - "46": 414038528.0, - "47": 414089216.0, - "48": 414139904.0, - "49": 414190592.0, - "50": 414241280.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": { @@ -284,4 +284,4 @@ "50": 0.05277 } } -} \ No newline at end of file +} diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py b/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py index 7cccbdc6aa0..9cb6036be5a 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py @@ -12,6 +12,7 @@ import megatron.core.parallel_state as mpu from megatron.core.distributed import DistributedDataParallelConfig +from megatron.core.distributed.finalize_model_grads import finalize_model_grads from megatron.core.distributed.fsdp.mcore_fsdp_adapter import FullyShardedDataParallel from megatron.core.distributed.fsdp.src.megatron_fsdp.mixed_precision import HAVE_TE_MXFP8TENSOR from megatron.core.fp8_utils import HAVE_TE @@ -1291,3 +1292,125 @@ def compare_losses(loss_a: float, loss_b: float, reference: str = "b"): better = "equal" return {"abs_diff": abs_diff, "rel_diff": rel_diff, "better": better} + + +class TestFsdpNonUnitBucketPreservation: + """Regression test for the non-FSDP-unit bucket preservation fix. + + Non-FSDP-unit buckets must remain persistently allocated: cross-module + readers (e.g. MambaMixer.forward, which reads self.conv1d.weight as a + raw tensor via causal_conv1d_fn(weight=...) instead of self.conv1d(...)) + can dereference them at any time. Before the fix, AllGatherPipeline.reset() + freed every bucket's gather scratch — including the non-unit one — and + because nn.Module.__call__ is never invoked on conv1d, no hook ever fires + to refresh conv1d.weight's wbuf.data (set_param_data=True) afterwards. + The next forward then dereferences freed storage. + + A small bucket_size is required so conv1d.weight lands in a different + non-unit bucket from MambaMixer's recurse=False params (A_log/D/dt_bias); + otherwise MambaMixer's own pre-forward hook gathers the shared bucket + and incidentally refreshes conv1d.weight, masking the bug. + """ + + @classmethod + def setup_class(cls): + Utils.initialize_model_parallel() + + @classmethod + def teardown_class(cls): + Utils.destroy_model_parallel() + + def test_non_unit_bucket_preserved_across_param_sync(self): + if not is_torch_min_version("2.4.0"): + pytest.skip("Megatron FSDP requires torch >= 2.4.0") + if Utils.world_size != 2: + pytest.skip("Requires exactly 2 GPUs (DP=2).") + pytest.importorskip("mamba_ssm") + pytest.importorskip("causal_conv1d") + pytest.importorskip("einops") + + from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec + from megatron.core.ssm.mamba_layer import MambaLayer + from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + from megatron.core.transformer.transformer_layer import TransformerLayer + + # MambaMixer's __init__ reads the 'model-parallel-rng' tracker. + model_parallel_cuda_manual_seed(0) + + # HIDDEN=256 is the floor: mamba's d_inner = 2*HIDDEN = 512, nheads = + # d_inner / mamba_head_dim(64) = 8, and nheads % mamba_num_groups(8) == 0. + HIDDEN = 256 + config = TransformerConfig( + num_layers=1, + hidden_size=HIDDEN, + num_attention_heads=4, + bf16=True, + params_dtype=torch.bfloat16, + normalization="RMSNorm", + attention_dropout=0.0, + hidden_dropout=0.0, + ) + + class HybridStack(torch.nn.Module): + def __init__(self, config, pg_collection): + super().__init__() + self.attn_layer = TransformerLayer( + config=config, + submodules=hybrid_stack_spec.submodules.attention_layer.submodules, + layer_number=1, + pg_collection=pg_collection, + add_layer_offset=False, + ) + self.mamba_layer = MambaLayer( + config=config, + submodules=hybrid_stack_spec.submodules.mamba_layer.submodules, + layer_number=2, + pg_collection=pg_collection, + ) + + def forward(self, hidden_states): + h, _ = self.attn_layer(hidden_states=hidden_states, attention_mask=None) + return self.mamba_layer(hidden_states=h) + + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=["tp", "cp"]) + model = HybridStack(config, pg_collection).cuda().to(torch.bfloat16) + + # bucket_size=4096 forces conv1d.weight into a separate non-unit bucket + # from MambaMixer's recurse=False params; fsdp_unit_modules=[TransformerLayer] + # matches the production default in mcore_fsdp_adapter.py. + fsdp_model = FullyShardedDataParallel( + config=config, + ddp_config=DistributedDataParallelConfig( + data_parallel_sharding_strategy="optim_grads_params", + overlap_grad_reduce=True, + overlap_param_gather=True, + bucket_size=4096, + use_megatron_fsdp=True, + ), + module=model, + fsdp_unit_modules=[TransformerLayer], + ) + + # Guard against future bucketing changes silently masking the bug. + non_unit_buckets = [ + pg + for pg in fsdp_model.param_and_grad_buffer.parameter_groups + if pg.fsdp_unit_id is None + ] + assert len(non_unit_buckets) >= 2, ( + f"Expected non-unit params to split across multiple buckets, " + f"got {len(non_unit_buckets)}. bucket_size may be too large." + ) + + x = torch.randn(64, 2, HIDDEN, device="cuda", dtype=torch.bfloat16) + out = fsdp_model(x) + out.sum().backward() + # Mirrors what the pipeline schedule does after backward: finalize_model_grads + # calls fsdp_model.finish_grad_sync(), which calls synchronize_param_gather() + # (because overlap_param_gather=True), which calls AllGatherPipeline.reset() — + # the buggy entry point that frees non-unit buckets. + finalize_model_grads([fsdp_model]) + fsdp_model(x) + + # Surface any deferred CUDA errors before teardown. + torch.cuda.synchronize()