Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / passes / pipeline.py: 98%
58 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""
16Pass Pipeline - Pass Orchestrator
18Automatically builds Pass execution order based on parallel configuration.
20Execution Order:
211. Basic optimization: DeadCodeElimination, CanonicalizeGraph
222. Activation checkpointing: ActivationCheckpointPass (``ac_mode`` != off)
23 runs on the plain joint graph BEFORE FSDPPass, so recomputed duplicates
24 are ordinary forward ops when the FSDP collectives are inserted (their
25 ``autograd_backward`` tag makes FSDPPass count them as backward
26 consumers and re-gather parameters for them). Running AC after FSDP
27 would make recompute read sharded placeholders.
282b. CPU offload: CpuOffloadPass (``activation_swap`` != none) moves saved
29 attention activations to host memory. It runs before FSDP so reload nodes
30 become ordinary backward consumers for parameter re-gathers. Attention swap
31 is incompatible with activation checkpointing (enforced in ``PassConfig``),
32 so it never runs alongside recompute.
333. Execution layer:
34 - FSDPPass (FSDP)
35 - FsdpBucketingPass (opt-in AG/RS fusion; supports FSDP with PP disabled).
36 - PpPass (pipeline-parallel stage split) — runs AFTER FSDPPass: the FSDP
37 collectives are value-level FX nodes that follow the data flow, so the
38 stage slices inherit this stage's all_gather/reduce_scatter untouched
39 while foreign stages' collectives are simply never copied. Splitting
40 first would hand FSDPPass two subgraphs and a loss-less output list,
41 breaking its grad-index contract.
424. Communication-compute overlap: AutoOverlapPass, then FsdpOverlapPass
43 (independent FSDP all-gather prefetch and gradient overlap).
445. Pipeline split: PpPass runs last among the graph-level passes, after
45 every graph-level pass. It replaces the joint graph with a schedule stub,
46 moving the collectives into the schedule's per-stage submodules where a
47 graph-level pass can no longer reach them. Running it last keeps the
48 non-PP pass order untouched and still lets FsdpOverlapPass reorder the
49 FSDP collectives before the split.
506. Backend compilation: InductorPass is TERMINAL (``inductor_mode`` != off).
51 It lowers the graph — or its tagged compute regions — to TorchInductor
52 kernels; afterwards the FX graph is a wrapper around the compiled callable
53 and no further FX-graph-level pass may run. When PpPass already split the
54 graph, the pass compiles each per-stage forward/backward subgraph owned by
55 the installed schedule instead of the (compute-less) stub.
56"""
58__all__ = ["PassPipeline"]
60from typing import TYPE_CHECKING, Any, List, Optional
62from ..pass_config import PassConfig
63from .activation_checkpoint import ActivationCheckpointPass
64from .base import GraphPass
65from .bucketing.bucketing_pass import FsdpBucketingPass
66from .cpu_offload import CpuOffloadPass
67from .inductor_pass import InductorPass
68from .overlap.schedule import AutoOverlapPass, FsdpOverlapPass
69from .parallel.fsdp_pass import FSDPPass
70from .parallel.pp_pass import PpPass
72if TYPE_CHECKING:
73 from torch import fx
74 from ..graph_parallel_plan import GraphParallelPlan
77class PassPipeline:
78 """Pass Pipeline Orchestrator.
80 Automatically builds Pass execution order based on parallel configuration.
81 """
83 def __init__(
84 self,
85 pass_config: PassConfig,
86 parallel_plan: Optional["GraphParallelPlan"] = None,
87 ) -> None:
88 """Initialize the pipeline with a parallel config and optional plan.
90 Args:
91 pass_config: Parallel configuration driving pass selection.
92 parallel_plan: Optional sharding plan forwarded to partitioning
93 passes (e.g. ``FSDPPass``).
94 """
95 self.config: PassConfig = pass_config
96 self.parallel_plan = parallel_plan
97 self.passes: List[GraphPass] = []
99 def build(self) -> "PassPipeline":
100 """Build Pass Pipeline."""
101 # 1. Basic optimization
102 self.passes.append(DeadCodeEliminationPass())
103 self.passes.append(CanonicalizeGraphPass())
105 # 2. Activation checkpointing: recompute duplicates must exist while
106 # the graph is still plain (no collectives yet), so FSDPPass inserts
107 # parameter all-gathers / backward re-gathers around them.
108 if self.config.ac_mode != "off":
109 self.passes.append(ActivationCheckpointPass())
111 # 2b. CPU offload: saved attention activations move to host memory.
112 # Runs before FSDP so reload nodes are ordinary backward consumers when
113 # parameter re-gathers are inserted. Attention swap is incompatible with
114 # activation checkpointing, so AC never runs alongside this pass.
115 if self.config.activation_swap != "none":
116 self.passes.append(CpuOffloadPass())
118 # 3. Execution layer: Parallel dimension partitioning
119 if getattr(self.config, "fsdp_enabled", False):
120 self.passes.append(FSDPPass(parallel_plan=self.parallel_plan))
121 if self.config.fsdp_bucketing:
122 self.passes.append(FsdpBucketingPass())
124 # 4. Communication-compute overlap optimization
125 if getattr(self.config, "enable_overlap", False):
126 self.passes.append(AutoOverlapPass())
127 if self.config.fsdp_allgather_prefetch or self.config.fsdp_overlap_grad_reduce:
128 self.passes.append(FsdpOverlapPass())
130 # 5. Pipeline split must run last: it replaces the joint graph with a
131 # schedule stub, moving the collectives into the schedule's per-stage
132 # submodules. Every graph-level pass (notably FsdpOverlapPass) must
133 # therefore see the collectives before the split.
134 if getattr(self.config, "pp_enabled", False):
135 self.passes.append(PpPass(parallel_plan=self.parallel_plan))
137 # 6. Terminal backend compilation: after this the FX graph wraps the
138 # compiled callable and no further FX-graph-level pass may run.
139 if self.config.inductor_mode != "off":
140 self.passes.append(InductorPass())
142 return self
144 def run(
145 self,
146 graph_module: "fx.GraphModule",
147 **kwargs: Any,
148 ) -> "fx.GraphModule":
149 """Execute all Passes against ``graph_module`` using the pipeline's config.
151 The pipeline is built from a single ``PassConfig``
152 (``from_config`` / ``__init__``); that same config drives every
153 pass. If a caller needs to run the same pipeline against a
154 different config, build a new pipeline.
156 Args:
157 graph_module: The FX GraphModule to transform.
158 **kwargs: Extra keyword arguments forwarded to every pass
159 (e.g. ``model`` / ``fsdp_group_name``). ``parallel_plan``
160 is auto-filled from the pipeline's plan when omitted.
162 Returns:
163 The transformed graph module.
164 """
165 # Ensure the parallel_plan is always available to passes that need
166 # it (FSDPPass reads it from kwargs), even if the caller did not
167 # pass it explicitly.
168 if "parallel_plan" not in kwargs and self.parallel_plan is not None:
169 kwargs["parallel_plan"] = self.parallel_plan
171 config = self.config
172 for graph_pass in self.passes:
173 graph_module = graph_pass.run(graph_module, config, **kwargs)
174 return graph_module
176 @classmethod
177 def from_config(
178 cls,
179 config: PassConfig,
180 parallel_plan: Optional["GraphParallelPlan"] = None,
181 ) -> "PassPipeline":
182 """Create Pipeline from configuration.
184 Args:
185 config: Parallel configuration driving pass selection.
186 parallel_plan: Optional sharding plan forwarded to partitioning
187 passes.
189 Returns:
190 The built pipeline, ready to ``run``.
191 """
192 return cls(config, parallel_plan).build()
195class DeadCodeEliminationPass(GraphPass):
196 """Dead code elimination Pass."""
198 name = "dead_code_elimination"
200 def run(
201 self,
202 graph_module: "fx.GraphModule",
203 pass_config: PassConfig,
204 **kwargs: Any,
205 ) -> "fx.GraphModule":
206 """Eliminate dead nodes and recompile the graph.
208 Args:
209 graph_module: The FX GraphModule to transform.
210 pass_config: Parallel configuration (unused by this pass).
211 **kwargs: Extra pass kwargs (unused).
213 Returns:
214 The transformed graph module.
215 """
216 graph_module.graph.eliminate_dead_code()
217 graph_module.recompile()
218 return graph_module
221class CanonicalizeGraphPass(GraphPass):
222 """Graph canonicalization Pass."""
224 name = "canonicalize_graph"
226 def run(
227 self,
228 graph_module: "fx.GraphModule",
229 pass_config: PassConfig,
230 **kwargs: Any,
231 ) -> "fx.GraphModule":
232 """Lint and recompile the graph to a canonical form.
234 Args:
235 graph_module: The FX GraphModule to transform.
236 pass_config: Parallel configuration (unused by this pass).
237 **kwargs: Extra pass kwargs (unused).
239 Returns:
240 The transformed graph module.
241 """
242 graph_module.graph.lint()
243 graph_module.recompile()
244 return graph_module