Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / passes / pipeline.py: 98%

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 

17 

18Automatically builds Pass execution order based on parallel configuration. 

19 

20Execution Order: 

211. Basic optimization: DeadCodeElimination, CanonicalizeGraph 

222. Execution layer: 

23 - FSDPPass (FSDP) 

24 - PpPass (pipeline-parallel stage split) — runs AFTER FSDPPass: the FSDP 

25 collectives are value-level FX nodes that follow the data flow, so the 

26 stage slices inherit this stage's all_gather/reduce_scatter untouched 

27 while foreign stages' collectives are simply never copied. Splitting 

28 first would hand FSDPPass two subgraphs and a loss-less output list, 

29 breaking its grad-index contract. 

303. Communication-compute overlap: AutoOverlapPass (a no-op under PP — the 

31 stage stub carries no wait_tensor nodes; the GPipe schedule overlaps 

32 comm/compute at the Python level) 

33 

34The backend compilation slot (InductorPass) is intentionally not wired yet; 

35add it here when an inductor backend integration lands. 

36""" 

37 

38__all__ = ["PassPipeline"] 

39 

40from typing import TYPE_CHECKING, Any, List, Optional 

41 

42from ..pass_config import PassConfig 

43from .base import GraphPass 

44from .overlap.schedule import AutoOverlapPass 

45from .parallel.fsdp_pass import FSDPPass 

46from .parallel.pp_pass import PpPass 

47 

48if TYPE_CHECKING: 

49 from torch import fx 

50 from ..graph_parallel_plan import GraphParallelPlan 

51 

52 

53class PassPipeline: 

54 """Pass Pipeline Orchestrator. 

55 

56 Automatically builds Pass execution order based on parallel configuration. 

57 """ 

58 

59 def __init__( 

60 self, 

61 pass_config: PassConfig, 

62 parallel_plan: Optional["GraphParallelPlan"] = None, 

63 ) -> None: 

64 """Initialize the pipeline with a parallel config and optional plan. 

65 

66 Args: 

67 pass_config: Parallel configuration driving pass selection. 

68 parallel_plan: Optional sharding plan forwarded to partitioning 

69 passes (e.g. ``FSDPPass``). 

70 """ 

71 self.config: PassConfig = pass_config 

72 self.parallel_plan = parallel_plan 

73 self.passes: List[GraphPass] = [] 

74 

75 def build(self) -> "PassPipeline": 

76 """Build Pass Pipeline.""" 

77 # 1. Basic optimization 

78 self.passes.append(DeadCodeEliminationPass()) 

79 self.passes.append(CanonicalizeGraphPass()) 

80 

81 # 2. Execution layer: Parallel dimension partitioning 

82 if getattr(self.config, "fsdp_enabled", False): 

83 self.passes.append(FSDPPass(parallel_plan=self.parallel_plan)) 

84 if getattr(self.config, "pp_enabled", False): 

85 self.passes.append(PpPass(parallel_plan=self.parallel_plan)) 

86 

87 # 3. Communication-compute overlap optimization 

88 if getattr(self.config, "enable_overlap", False): 

89 self.passes.append(AutoOverlapPass()) 

90 

91 return self 

92 

93 def run( 

94 self, 

95 graph_module: "fx.GraphModule", 

96 **kwargs: Any, 

97 ) -> "fx.GraphModule": 

98 """Execute all Passes against ``graph_module`` using the pipeline's config. 

99 

100 The pipeline is built from a single ``PassConfig`` 

101 (``from_config`` / ``__init__``); that same config drives every 

102 pass. If a caller needs to run the same pipeline against a 

103 different config, build a new pipeline. 

104 

105 Args: 

106 graph_module: The FX GraphModule to transform. 

107 **kwargs: Extra keyword arguments forwarded to every pass 

108 (e.g. ``model`` / ``fsdp_group_name``). ``parallel_plan`` 

109 is auto-filled from the pipeline's plan when omitted. 

110 

111 Returns: 

112 The transformed graph module. 

113 """ 

114 # Ensure the parallel_plan is always available to passes that need 

115 # it (FSDPPass reads it from kwargs), even if the caller did not 

116 # pass it explicitly. 

117 if "parallel_plan" not in kwargs and self.parallel_plan is not None: 

118 kwargs["parallel_plan"] = self.parallel_plan 

119 

120 config = self.config 

121 for graph_pass in self.passes: 

122 graph_module = graph_pass.run(graph_module, config, **kwargs) 

123 return graph_module 

124 

125 @classmethod 

126 def from_config( 

127 cls, 

128 config: PassConfig, 

129 parallel_plan: Optional["GraphParallelPlan"] = None, 

130 ) -> "PassPipeline": 

131 """Create Pipeline from configuration.""" 

132 return cls(config, parallel_plan).build() 

133 

134 

135class DeadCodeEliminationPass(GraphPass): 

136 """Dead code elimination Pass.""" 

137 

138 name = "dead_code_elimination" 

139 

140 def run( 

141 self, 

142 graph_module: "fx.GraphModule", 

143 pass_config: PassConfig, 

144 **kwargs: Any, 

145 ) -> "fx.GraphModule": 

146 """Eliminate dead nodes and recompile the graph.""" 

147 graph_module.graph.eliminate_dead_code() 

148 graph_module.recompile() 

149 return graph_module 

150 

151 

152class CanonicalizeGraphPass(GraphPass): 

153 """Graph canonicalization Pass.""" 

154 

155 name = "canonicalize_graph" 

156 

157 def run( 

158 self, 

159 graph_module: "fx.GraphModule", 

160 pass_config: PassConfig, 

161 **kwargs: Any, 

162 ) -> "fx.GraphModule": 

163 """Lint and recompile the graph to a canonical form.""" 

164 graph_module.graph.lint() 

165 graph_module.recompile() 

166 return graph_module