add manual bucketing pass

Author: ruisizhang123

torchtitan/experiments/simple_fsdp/README.md

@@ -51,13 +51,14 @@ SimpleFSDP relies on compiler backend to perform optimizations (i.e., bucketing
 
 2. auto optimization: perform auto-bucketing & reordering without user inputs. **Note: it is not guaranteed that users will get the most optimized training performance**
     - "aot_eager_autobucketing": perform autobucketing at aten fx-level, and perform code execution with aot_eager backend.
-
-
-users can specify the pass (e.g., "aot_eager_autobucketing") via additional configs:
-
-```bash
---job.custom_config_module=torchtitan.experiments.simple_fsdp.job_config  --compile.model_backend_override "aot_eager_autobucketing"
-```
+      ```bash
+      --compile.backend "aot_eager" --job.custom_config_module=torchtitan.experiments.simple_fsdp.job_config  --compile.model_backend_override "aot_eager_autobucketing"
+      ```
+3. manual optimization: perform manual bucketing & reordering with user FQN inputs.
+    - "aot_eager_manualbucketing": perform manual bucketing at aten fx-level, and perform code execution with aot_eager backend.
+      ```bash
+      --compile.backend "aot_eager" --job.custom_config_module=torchtitan.experiments.simple_fsdp.job_config  --compile.model_backend_override "aot_eager_manualbucketing" --compile.manual_bucketed_modules "tok_embeddings,layers.[0-5],norm+output"
+      ```
 
 ### Citation
 

torchtitan/experiments/simple_fsdp/backend.py

@@ -7,12 +7,18 @@
 from typing import Any, Union
 
 import torch
+from ... import Compile as CompileConfig
 
 
-def get_compile_backend(backend_name: str) -> Union[str, callable]:
+def get_compile_backend(compile_config: CompileConfig) -> Union[str, callable]:
     # return the compile backends used in SimpleFSDP training
     # Step1: check if backend_name is inside available torch.compile backends
     # Step2: check if the backend_name has been registered as a customized backend
+    backend_name = (
+        getattr(compile_config, "model_backend_override", None)
+        or compile_config.backend
+    )
+
     available_torch_backend = torch._dynamo.list_backends(exclude_tags=())
     if backend_name in available_torch_backend:
         return backend_name
@@ -25,8 +31,6 @@ def get_compile_backend(backend_name: str) -> Union[str, callable]:
             schedule_overlap_bucketing,
         )
 
-        torch._inductor.config.test_configs.aten_fx_overlap_preserving_bucketing = True
-        torch._inductor.config.test_configs.aten_fx_overlap_insert_overlap_deps = False
         torch._inductor.config.allow_buffer_reuse = False
 
         def aten_autobucketing_reordering_pass(
@@ -41,6 +45,34 @@ def aten_autobucketing_reordering_pass(
             bw_compiler=aten_autobucketing_reordering_pass,
             keep_inference_input_mutations=True,
         )
+    elif backend_name == "aot_eager_manualbucketing":
+        # Perform manual optimization in aten fx-level and execute code in aot_eager backend
+        # The manualbucketing logic is here:
+        bucketing_modules = compile_config.manual_bucketed_modules
+        from functools import partial
+
+        from torch._dynamo.backends.common import aot_autograd as aot_autograd_backend
+        from torch._inductor.fx_passes.overlap_manual_scheduling import (
+            manual_overlap_bucketing,
+        )
+
+        torch._inductor.config.allow_buffer_reuse = False
+        manual_overlap_bucketing = partial(
+            manual_overlap_bucketing,
+            module_bucket_plans=compile_config.manual_bucketed_modules,
+        )
+
+        def aten_manualbucketing_reordering_pass(
+            gm: torch.fx.GraphModule, example_inputs: Any
+        ) -> torch.fx.GraphModule:
+            manual_overlap_bucketing(gm)
+            return gm
+
+        backend = aot_autograd_backend(
+            fw_compiler=aten_manualbucketing_reordering_pass,
+            bw_compiler=aten_manualbucketing_reordering_pass,
+            keep_inference_input_mutations=True,
+        )
     else:
         raise AssertionError(f"Unsupported customized backend: {backend_name}")
 

torchtitan/experiments/simple_fsdp/deepseek_v3/parallelize.py

@@ -162,6 +162,6 @@ def parallelize_deepseekv3(
     if job_config.compile.enable:
         torch._inductor.config.reorder_for_peak_memory = False
         torch._dynamo.config.capture_scalar_outputs = True
-        model = torch.compile(model, backend=job_config.compile.backend, fullgraph=True)
+        model = torch.compile(model, backend=get_compile_backend(job_config.compile), fullgraph=True)
 
     return model

torchtitan/experiments/simple_fsdp/job_config.py

@@ -10,7 +10,18 @@
 @dataclass
 class Compile:
     model_backend_override: str | None = None
-    """Override backend to compile in simplefsdp. Additional backend includes aot_eager_autobucketing"""
+    """Override backend to compile in simplefsdp. Additional backend includes aot_eager_autobucketing """
+
+    manual_bucketed_modules: list[str] = field(default_factory=list)
+    """
+    Manual bucket modules based on user specified FQNs
+    Abbreviations are supported to make specifying modules easier.
+    Currently, the following abbreviations are available:
+    (1) layers.[0-2] -> [layers.0], [layers.1], [layers.2]
+        (layers are split three separate buckets)
+    (2) norm+output -> [norm, output]
+        (norm and output are in one bucket)
+    """
 
 
 @dataclass

torchtitan/experiments/simple_fsdp/llama3/parallelize.py

@@ -125,13 +125,9 @@ def parallelize_llama(
 
     if job_config.compile.enable and "model" in job_config.compile.components:
         torch._inductor.config.reorder_for_peak_memory = False
-        backend = (
-            getattr(job_config.compile, "model_backend_override", None)
-            or job_config.compile.backend
-        )
         model = torch.compile(
             model,
-            backend=get_compile_backend(backend),
+            backend=get_compile_backend(job_config.compile),
             fullgraph=True,
         )
 

torchtitan/experiments/simple_fsdp/tests/integration_tests.py

@@ -29,18 +29,19 @@ def build_simple_fsdp_test_list() -> list[OverrideDefinitions]:
             "1D",
             "1d",
         ),
-        OverrideDefinitions(
-            [
-                [
-                    "--model.name simple_fsdp.llama3",
-                    "--compile.enable",
-                    "--job.custom_config_module=torchtitan.experiments.simple_fsdp.job_config",
-                    "--compile.model_backend_override aot_eager_autobucketing",
-                ],
-            ],
-            "1D+aot_eager_autobucketing",
-            "1d_aot_eager_autobucketing",
-        ),
+        # TODO(ruisizhang123): add back after autobucketing pass is mature
+        # OverrideDefinitions(
+        #     [
+        #         [
+        #             "--model.name simple_fsdp.llama3",
+        #             "--compile.enable",
+        #             "--job.custom_config_module=torchtitan.experiments.simple_fsdp.job_config",
+        #             "--compile.model_backend_override aot_eager_autobucketing",
+        #         ],
+        #     ],
+        #     "1D+aot_eager_autobucketing",
+        #     "1d_aot_eager_autobucketing",
+        # ),
         OverrideDefinitions(
             [
                 [