Make Megatron-FSDP torch.compile compatible

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:

megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py

@@ -611,6 +611,7 @@ def _post_backward(module, *unused):
                     ),
                 )
 
+        @torch.compiler.disable
         def _pre_forward_param_unshard(
             module: nn.Module, args: Tuple[Any, ...], kwargs: Dict[str, Any]
         ):
@@ -642,6 +643,7 @@ def _pre_forward_param_unshard(
             )
             return args, kwargs
 
+        @torch.compiler.disable
         def _register_post_backward_hook(
             post_backward_hook: callable,
             module: nn.Module,
@@ -733,6 +735,7 @@ def _root_post_backward(*unused):
             if self.model_auto_sync:
                 self.finish_grad_sync()
 
+        @torch.compiler.disable
         def _pre_backward_param_unshard(module: nn.Module, *unused):
             """
             Sub-module pre-backward hook to all-gather the module parameters
@@ -796,6 +799,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
@@ -815,6 +819,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))
 
@@ -824,6 +829,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

@@ -232,16 +232,23 @@ 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(
+        "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,
+        self, model_type, dp_shard_strategy, dp_outer_strategy, mesh_dim_config, common_args
     ):
         """
         Test the fully_shard API with different configurations.
@@ -253,6 +260,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(
@@ -297,6 +308,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")