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 torchtitan.config import JobConfig
 
 
-def get_compile_backend(backend_name: str) -> Union[str, callable]:
+def get_compile_backend(job_config: JobConfig) -> 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(job_config.compile, "model_backend_override", None)
+        or job_config.compile.backend
+    )
+
     available_torch_backend = torch._dynamo.list_backends(exclude_tags=())
     if backend_name in available_torch_backend:
         return backend_name
@@ -41,6 +47,36 @@ 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 = job_config.compile.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.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
+        manual_overlap_bucketing = partial(
+            manual_overlap_bucketing,
+            module_bucket_plans=job_config.compile.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/job_config.py

@@ -10,7 +10,10 @@
 @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)
+    """Which modules should be bucketed together based on user specifications in manual optimization """
 
 
 @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),
             fullgraph=True,
         )