Make Megatron-FSDP torch.compile compatible

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

tests/unit_tests/distributed/fsdp/test_mfsdp_fully_shard.py

@@ -235,6 +235,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,
@@ -243,6 +244,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.
@@ -299,6 +301,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")