From f89d9d615558b0f8908ed2be73cd318d31edd83e Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Thu, 14 May 2026 22:19:08 +0000 Subject: [PATCH 1/2] Use StorageResizeBasedBucketAllocator as FixedPoolAllocator backup When fsdp_double_buffer=True, FSDP units whose bucket structure does not match the rep-unit fall through to FixedPoolAllocator's backup allocator. The previous backup, TemporaryBucketAllocator, discards the Bucket (and its Storage object) on free and constructs a fresh torch.empty on the next allocate. Autograd's SavedVariable for non-leaf views (e.g. rearrange(conv1d.weight) saved by causal_conv1d_fn) holds a reference to the original Storage object; after the discard, that reference is permanently empty, and a backward dispatch through the saved view yields a CUDA illegal memory access. Switch the backup to StorageResizeBasedBucketAllocator, which retains the Bucket/Storage object across free/allocate cycles via _resize_(0) and _resize_(size). The saved view's storage handle stays valid; the kernel reads through it to whatever address the pre-bwd regather installed. This matches the allocator already used as the default for the non-double-buffer path. Add TestFsdpHybridModelDoubleBuffer to exercise the path with an attention + Mamba stack. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../megatron_fsdp/param_and_grad_buffer.py | 2 +- .../test_mcore_fully_sharded_data_parallel.py | 111 ++++++++++++++++++ 2 files changed, 112 insertions(+), 1 deletion(-) 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 266da6b74c4..52140580916 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 @@ -724,7 +724,7 @@ def __init__( # Fallback allocator used if the fixed pool allocator cannot fulfill a request. self.fallback_to_persistent_buffer = fallback_to_persistent_buffer - self.backup_allocator = TemporaryBucketAllocator() + self.backup_allocator = StorageResizeBasedBucketAllocator() def _is_two_bucket_group_equal(self, group_a, group_b): # Check if two bucket groups are equivalent in dtype and size. 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 df5eac3659a..eacbb240286 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.hyper_comm_grid import HyperCommGrid @@ -1070,3 +1071,113 @@ 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 TestFsdpHybridModelDoubleBuffer: + """Smoke test: hybrid (attention + Mamba) model trained for a few steps + under Megatron FSDP with TransformerLayer and MambaLayer marked as FSDP + unit modules and fsdp_double_buffer=True. + """ + + @classmethod + def setup_class(cls): + Utils.initialize_model_parallel() + + @classmethod + def teardown_class(cls): + Utils.destroy_model_parallel() + + def test_train_steps_with_double_buffer(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 + + model_parallel_cuda_manual_seed(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.transformer_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.transformer_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) + + 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, + fsdp_double_buffer=True, + ), + module=model, + fsdp_unit_modules=[TransformerLayer, MambaLayer], + ) + + optimizer = DistributedOptimizer( + optimizer=None, + config=OptimizerConfig(optimizer="adam", lr=1e-3), + grad_scaler=None, + init_state_fn=None, + model_chunks=[fsdp_model], + per_model_buffers={0: [fsdp_model.param_and_grad_buffer]}, + data_parallel_group=fsdp_model.megatron_fsdp_dist_index.get_dp_group(), + data_parallel_group_gloo=None, + data_parallel_group_idx=0, + distributed_optimizer_instance_id=0, + ) + + NUM_MICROBATCHES = 4 + for _ in range(5): + optimizer.zero_grad() + for microbatch_idx in range(NUM_MICROBATCHES): + fsdp_model.is_last_microbatch = microbatch_idx == NUM_MICROBATCHES - 1 + x = torch.randn(64, 2, HIDDEN, device="cuda", dtype=torch.bfloat16) + out = fsdp_model(x) + loss = out.sum() + loss.backward() + finalize_model_grads([fsdp_model]) + optimizer.step() + + torch.cuda.synchronize() From b6159a0bb6ddaebafcafb7bd59e0e0e9decb4b39 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Fri, 15 May 2026 03:25:13 +0000 Subject: [PATCH 2/2] Fix lint: collapse use_mpu_process_groups call onto one line Co-Authored-By: Claude Opus 4.7 (1M context) --- .../megatron_fsdp/test_mcore_fully_sharded_data_parallel.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) 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 eacbb240286..4b2c6ab80cf 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 @@ -1136,9 +1136,7 @@ def forward(self, hidden_states): h, _ = self.transformer_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"] - ) + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=["tp", "cp"]) model = HybridStack(config, pg_collection).cuda().to(torch.bfloat16) fsdp_model = FullyShardedDataParallel(