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: 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 e2cbccf4356..d93a13d241b 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/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 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 b0bd6c729ef..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,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")