77import torch
88import torch .nn as nn
99from torch ._subclasses .fake_tensor import FakeTensorMode
10- from torch .distributed .algorithms ._checkpoint .checkpoint_wrapper import (
11- CheckpointWrapper ,
12- )
1310
1411from torchtitan .config import CompileConfig
15- from torchtitan .models .common import moe as moe_module
1612from torchtitan .tools .logging import logger
1713
1814
2420)
2521
2622
27- def apply_compile_dense (model : nn .Module , compile_config : CompileConfig ) -> None :
28- """
29- Apply torch.compile to each TransformerBlock, which makes compilation efficient due to
30- repeated structure. Alternatively one can compile the whole model (after applying DP).
31-
32- This is for dense (non-MoE) models. It compiles each TransformerBlock as a whole.
33- """
34- # Skip replaying forward side effects (e.g. RoPE cache updates) during
35- # the AC recompute in backward. Eager AC replays the forward python
36- # side-effects in backward, but torch.compile has no easy way to reapply
37- # python mutations in the backward. Setting this flag accepts this eager
38- # and compile divergence by skipping reapplication of side effects.
39- torch ._dynamo .config .skip_fwd_side_effects_in_bwd_under_checkpoint = (
40- True # pyrefly: ignore [bad-assignment]
41- )
42-
43- # pyrefly: ignore [missing-attribute]
44- for layer_id , transformer_block in model .layers .named_children ():
45- transformer_block .compile (backend = compile_config .backend , fullgraph = True )
46- # pyrefly: ignore [missing-attribute]
47- model .layers .register_module (layer_id , transformer_block )
48-
49- logger .info ("Compiling each TransformerBlock with torch.compile" )
50-
51-
52- def apply_compile_sparse (
53- model : nn .Module , compile_config : CompileConfig , ep_enabled : bool
54- ) -> None :
23+ def apply_compile (model : nn .Module , compile_config : CompileConfig ) -> None :
5524 """
5625 Apply torch.compile to each TransformerBlock, which makes compilation efficient due to
5726 repeated structure. Alternatively one can compile the whole model (after applying DP).
58-
59- This is for MoE (sparse) models. It compiles sub-modules individually to avoid
60- graph breaks from FSDP(GroupedExperts).
6127 """
62- # Needed for torch.compile to avoid graph breaking on dynamic shapes in
63- # token-choice MoE, but it is experimental .
28+ # Needed for torch.compile to handle data-dependent dynamic shapes in
29+ # token-choice MoE dispatch. Harmless for dense models .
6430 torch ._dynamo .config .capture_scalar_outputs = True
6531 # Skip replaying forward side effects (e.g. RoPE cache updates) during
6632 # the AC recompute in backward. Eager AC replays the forward python
@@ -73,77 +39,6 @@ def apply_compile_sparse(
7339
7440 # pyrefly: ignore [missing-attribute]
7541 for layer_id , transformer_block in model .layers .named_children ():
76- if transformer_block .moe_enabled :
77- # If it is a MoE layer, FSDP(GroupedExperts) will cause a graph break
78- # So we must weave compile wrappers around those FSDP hooks to
79- # prevent AC from falling back the whole graph to eager.
80- # TODO: Fix Compile(AC(graph break))
81-
82- if isinstance (transformer_block , CheckpointWrapper ):
83- # TODO: Make CheckpointWrapper a transparent wrapper
84- # unwrap so that .named_children() works
85- block = transformer_block ._checkpoint_wrapped_module
86- else :
87- block = transformer_block
88-
89- for attr_name , submod in block .named_children ():
90- assert getattr (block , attr_name ) == getattr (
91- transformer_block , attr_name
92- )
93-
94- if isinstance (submod , moe_module .MoE ):
95- # avoid graph breaking on the GroupedExperts' FSDP hooks
96- # by wrapping each submod's forward instead of their __call__
97- moe = submod
98- for attr_name , submod in moe .named_children ():
99- if attr_name == "experts" :
100- # NOTE: We don't compile token dispatch and token combine due to an issue on B200:
101- # https://github.com/pytorch/torchtitan/issues/1940
102- continue
103- submod .compile (backend = compile_config .backend , fullgraph = True )
104- else :
105- submod .compile (backend = compile_config .backend , fullgraph = True )
106- else :
107- # If it's not a MoE layer, there is no FSDP(GroupedExperts)
108- # So we can compile the whole block
109- transformer_block .compile (
110- backend = compile_config .backend ,
111- fullgraph = True ,
112- )
113-
114- # pyrefly: ignore [missing-attribute]
115- model .layers .register_module (layer_id , transformer_block )
116-
117- # Patch some globals only once (apply_compile_sparse is called multiple times for PP setup)
118- already_patched = (
119- "_run_experts_grouped_mm_dynamic"
120- in moe_module ._run_experts_grouped_mm .__qualname__
121- )
122- if not already_patched :
123- moe_module ._run_experts_grouped_mm = torch .compile (
124- moe_module ._run_experts_grouped_mm ,
125- backend = compile_config .backend ,
126- fullgraph = True ,
127- )
128-
129- if ep_enabled :
130- compiled_fn = moe_module ._run_experts_grouped_mm
131-
132- # keep function logic in sync with `already_patched` above
133- def _run_experts_grouped_mm_dynamic (
134- w1 : torch .Tensor ,
135- w2 : torch .Tensor ,
136- w3 : torch .Tensor ,
137- x : torch .Tensor ,
138- num_tokens_per_expert : torch .Tensor ,
139- ) -> torch .Tensor :
140- # dynamic number of tokens in expert parallel
141- torch ._dynamo .mark_dynamic (x , 0 )
142- return compiled_fn (w1 , w2 , w3 , x , num_tokens_per_expert )
143-
144- moe_module ._run_experts_grouped_mm = _run_experts_grouped_mm_dynamic
145-
146- # NOTE: We don't compile for loop code path due to an issue with unbacked symints:
147- # https://github.com/pytorch/pytorch/issues/166460
42+ transformer_block .compile (backend = compile_config .backend , fullgraph = True )
14843
14944 logger .info ("Compiling each TransformerBlock with torch.compile" )
0 commit comments