From 26932de60debec8057e4eca4b756a277866f8103 Mon Sep 17 00:00:00 2001 From: jianbinc Date: Fri, 28 Nov 2025 19:25:05 +0800 Subject: [PATCH 1/5] disable m-fsdp hooks for torch compile compatible --- .../distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py | 6 ++++++ 1 file changed, 6 insertions(+) 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 8a63e0f5cf7..5b340017c21 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -595,6 +595,7 @@ def _post_backward(module, *unused): ), ) + @torch.compiler.disable def _pre_forward_param_unshard( module: nn.Module, args: Tuple[Any, ...], kwargs: Dict[str, Any] ): @@ -623,6 +624,7 @@ def _pre_forward_param_unshard( ) return args, kwargs + @torch.compiler.disable def _register_post_backward_hook( post_backward_hook: callable, module: nn.Module, @@ -714,6 +716,7 @@ def _root_post_backward(*unused): if self.model_auto_sync: self.finish_grad_sync() + @torch.compiler.disable def _pre_backward(module: nn.Module, *unused): """ Sub-module pre-backward hook to all-gather the module parameters @@ -767,6 +770,7 @@ def _root_pre_backward(module: nn.Module, *unused): # the backward pass. torch.autograd.Variable._execution_engine.queue_callback(_root_post_backward) + @torch.compiler.disable def _post_forward(module: nn.Module, input: Any, output: Any): # When composed with module-hook-based activation recomputation, the # post-backward hook is responsible for resharding the module parameters @@ -782,6 +786,7 @@ def _post_forward(module: nn.Module, input: Any, output: Any): return output + @torch.compiler.disable def _release_module_fp8_transpose_cache(module: nn.Module, *unused): release_params_fp8_transpose_cache(module.parameters(recurse=False)) @@ -791,6 +796,7 @@ def create_custom_backward_hook(module, custom_backward_handler): to the output tensor(s) of a module during a post-forward hook. """ + @torch.compiler.disable def forward_hook(_module, inputs, output): # Replace the output to avoid the output tensor being the same as # the input tensor, which makes it impossible to identify which From f45c937bc7cfd04b07ec5e045a4fbdaf7187d661 Mon Sep 17 00:00:00 2001 From: jianbinc Date: Tue, 16 Dec 2025 23:54:04 +0800 Subject: [PATCH 2/5] add torch compile in mfsdp UT --- tests/unit_tests/distributed/fsdp/test_mfsdp_fully_shard.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unit_tests/distributed/fsdp/test_mfsdp_fully_shard.py b/tests/unit_tests/distributed/fsdp/test_mfsdp_fully_shard.py index bbf2c8f7689..e0fce82bf70 100644 --- a/tests/unit_tests/distributed/fsdp/test_mfsdp_fully_shard.py +++ b/tests/unit_tests/distributed/fsdp/test_mfsdp_fully_shard.py @@ -225,6 +225,7 @@ def teardown_class(cls): ) @pytest.mark.parametrize("preserve_fp32_weights", [True, False]) @pytest.mark.parametrize("init_model_with_meta_device", [True, False]) + @pytest.mark.parametrize("torch_compile", [True, False]) def test_fully_shard( self, model_type, @@ -233,6 +234,7 @@ def test_fully_shard( mesh_dim_config, preserve_fp32_weights, init_model_with_meta_device, + torch_compile, ): """ Test the fully_shard API with different configurations. @@ -289,6 +291,7 @@ def test_fully_shard( grad_reduce_in_fp32=False, init_model_with_meta_device=init_model_with_meta_device, ) + model = torch.compile(model) if torch_compile else model # Mock input and target. toy_input = torch.randn(1, DIM_SIZE, DIM_SIZE).to("cuda") From 4f1ddb8afe1f373149400b669975a38efd36301a Mon Sep 17 00:00:00 2001 From: jianbinc Date: Sat, 10 Jan 2026 23:34:16 +0800 Subject: [PATCH 3/5] Add torch.compile document in M-FSDP README --- megatron/core/distributed/fsdp/src/README.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/megatron/core/distributed/fsdp/src/README.md b/megatron/core/distributed/fsdp/src/README.md index 53ac1ff81b7..c96e4425f9a 100644 --- a/megatron/core/distributed/fsdp/src/README.md +++ b/megatron/core/distributed/fsdp/src/README.md @@ -116,6 +116,10 @@ fully_shard(model) # Your model is now ready for distributed training! ``` +### `torch.compile` support + +Megatron-FSDP supports `torch.compile`, but this feature is still experimental and may introduce performance regressions in some workloads. + ## `fully_shard` / `MegatronFSDP` API - Advanced Features Megatron-FSDP's `fully_shard_*` API has a comprehensive set of arguments for fine-tuning your model's performance: From b57a671695008434cde38a4f5425050b082dd138 Mon Sep 17 00:00:00 2001 From: jianbinc Date: Mon, 12 Jan 2026 14:46:12 +0800 Subject: [PATCH 4/5] Simplify test_mfsdp_fully_shard.py test cases --- .../fsdp/test_mfsdp_fully_shard.py | 23 ++++++++++++++----- 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/tests/unit_tests/distributed/fsdp/test_mfsdp_fully_shard.py b/tests/unit_tests/distributed/fsdp/test_mfsdp_fully_shard.py index 6961ba911d3..8f702bd0669 100644 --- a/tests/unit_tests/distributed/fsdp/test_mfsdp_fully_shard.py +++ b/tests/unit_tests/distributed/fsdp/test_mfsdp_fully_shard.py @@ -232,18 +232,25 @@ def teardown_class(cls): (2, 2, 1, 2), ], ) - @pytest.mark.parametrize("preserve_fp32_weights", [True, False]) - @pytest.mark.parametrize("init_model_with_meta_device", [True, False]) - @pytest.mark.parametrize("torch_compile", [True, False]) + @pytest.mark.parametrize("common_args", [ + { + "preserve_fp32_weights": True, + "init_model_with_meta_device": True, + "torch_compile": True + }, + { + "preserve_fp32_weights": False, + "init_model_with_meta_device": False, + "torch_compile": False + } + ]) def test_fully_shard( self, model_type, dp_shard_strategy, dp_outer_strategy, mesh_dim_config, - preserve_fp32_weights, - init_model_with_meta_device, - torch_compile, + common_args, ): """ Test the fully_shard API with different configurations. @@ -255,6 +262,10 @@ def test_fully_shard( """ from megatron.core.distributed.fsdp.src.megatron_fsdp.fully_shard import fully_shard + preserve_fp32_weights = common_args["preserve_fp32_weights"] + init_model_with_meta_device = common_args["init_model_with_meta_device"] + torch_compile = common_args["torch_compile"] + # Skip due to lack of functionality. if init_model_with_meta_device and dp_shard_strategy == NO_SHARD: pytest.skip( From d2260d7b90df6e3c08b6f78695d6ef2586a7b410 Mon Sep 17 00:00:00 2001 From: jianbinc Date: Mon, 12 Jan 2026 14:50:21 +0800 Subject: [PATCH 5/5] code format --- .../fsdp/test_mfsdp_fully_shard.py | 34 +++++++++---------- 1 file changed, 16 insertions(+), 18 deletions(-) diff --git a/tests/unit_tests/distributed/fsdp/test_mfsdp_fully_shard.py b/tests/unit_tests/distributed/fsdp/test_mfsdp_fully_shard.py index 8f702bd0669..191aac3e01b 100644 --- a/tests/unit_tests/distributed/fsdp/test_mfsdp_fully_shard.py +++ b/tests/unit_tests/distributed/fsdp/test_mfsdp_fully_shard.py @@ -232,25 +232,23 @@ def teardown_class(cls): (2, 2, 1, 2), ], ) - @pytest.mark.parametrize("common_args", [ - { - "preserve_fp32_weights": True, - "init_model_with_meta_device": True, - "torch_compile": True - }, - { - "preserve_fp32_weights": False, - "init_model_with_meta_device": False, - "torch_compile": False - } - ]) + @pytest.mark.parametrize( + "common_args", + [ + { + "preserve_fp32_weights": True, + "init_model_with_meta_device": True, + "torch_compile": True, + }, + { + "preserve_fp32_weights": False, + "init_model_with_meta_device": False, + "torch_compile": False, + }, + ], + ) def test_fully_shard( - self, - model_type, - dp_shard_strategy, - dp_outer_strategy, - mesh_dim_config, - common_args, + self, model_type, dp_shard_strategy, dp_outer_strategy, mesh_dim_config, common_args ): """ Test the fully_shard API with different configurations.