Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_ppb / pp_optimizer.py: 80%

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"""PP optimizer — YAML + JSON driven pipeline parallelism strategy optimizer.""" 

16 

17from __future__ import annotations 

18 

19import logging 

20import math 

21from typing import Any 

22 

23from hyper_parallel.auto_parallel.sapp_ppb.pp_config_builder.layer_loader import ( 

24 LayerBuilder, 

25) 

26from hyper_parallel.auto_parallel.sapp_ppb.pp_modeling.pp_structs import ( 

27 PPStrategyResult, 

28) 

29from hyper_parallel.auto_parallel.sapp_ppb.pp_modeling.pp_balancer import ( 

30 PPBalancer, 

31) 

32from hyper_parallel.auto_parallel.sapp_ppb.pp_config_builder.yaml_parser import ( 

33 parse_yaml_for_optimization, 

34) 

35from hyper_parallel.auto_parallel.sapp_ppb.pp_sim_adapter import PPSimulator 

36 

37 

38class PPOptimizer: 

39 """Unified PP strategy optimizer — YAML + JSON driven. 

40 

41 Pipeline topology and ILP constraints (pp_degree, num_layer, 

42 micro_batch_num, num_of_interleave, memory_limit, etc.) come from 

43 a YAML file; layer descriptions come from a JSON file. 

44 

45 Orchestration flow: 

46 

47 1. Parse YAML config and build layers from JSON. 

48 2. Call :class:`PPBalancer` to solve the ILP and build the 

49 strategy result. 

50 3. If ``enable_simulation`` is ``True``, run :class:`PPSimulator` 

51 and merge simulation metrics into the result. 

52 

53 Example: 

54 >>> optimizer = PPOptimizer() 

55 >>> result = optimizer.optimize( 

56 ... yaml_path="init_demo.yaml", 

57 ... json_path="demo.json", 

58 ... ) 

59 """ 

60 

61 def optimize( 

62 self, 

63 yaml_path: str = "", 

64 json_path: str = "", 

65 ) -> PPStrategyResult: 

66 """Optimize PP strategy — YAML + JSON driven. 

67 

68 Reads pipeline topology and ILP constraints from ``yaml_path`` 

69 and layer descriptions from ``json_path``, then runs a single 

70 ILP optimisation, optionally runs the pipeline simulator, and 

71 returns the result. 

72 

73 Args: 

74 yaml_path: Path to the YAML configuration file with pipeline 

75 topology and ILP constraints (``pipeline_num``, 

76 ``num_layer``, ``micro_batch_num``, 

77 ``num_of_interleave``, ``memory_limit``, etc.). 

78 json_path: Path to the JSON file with ``layers_description`` 

79 defining the model layers (HEAD, BODY, TAIL). 

80 

81 Returns: 

82 Optimised PP strategy result. 

83 

84 Raises: 

85 ValueError: If ``yaml_path`` or ``json_path`` is empty. 

86 RuntimeError: If the ILP solution is infeasible. 

87 

88 Example: 

89 >>> optimizer = PPOptimizer() 

90 >>> result = optimizer.optimize( 

91 ... yaml_path="init_demo.yaml", 

92 ... json_path="demo.json", 

93 ... ) 

94 """ 

95 if not yaml_path: 

96 raise ValueError( 

97 "PPOptimizer.optimize requires yaml_path. " 

98 "Please provide a valid YAML configuration path." 

99 ) 

100 if not json_path: 

101 raise ValueError( 

102 "PPOptimizer.optimize requires json_path. " 

103 "Please provide a valid JSON profile path." 

104 ) 

105 

106 yaml_config = parse_yaml_for_optimization(yaml_path) 

107 

108 builder = LayerBuilder(yaml_config, json_path) 

109 balancer = PPBalancer(builder) 

110 result = balancer.balance_with_ilp() 

111 

112 if not result.is_feasible: 

113 details = result.infeasibility_details 

114 msg = f"PP optimization failed: {details.get('reason', 'unknown')}" 

115 if details.get("error"): 

116 msg += f" ({details['error']})" 

117 raise RuntimeError(msg) 

118 

119 if yaml_config.enable_simulation: 

120 self._run_simulation( 

121 result, 

122 balancer.pipeline, 

123 yaml_config, 

124 builder.constant_memory, 

125 yaml_config.sim_comm_time, 

126 ) 

127 

128 if result.simulation_status == "failed": 

129 logging.getLogger(__name__).warning( 

130 "ILP succeeded but simulation failed: %s. " 

131 "Pipeline bubble and step time estimates are unavailable.", 

132 result.simulation_error, 

133 ) 

134 elif result.simulation_status == "not_run": 

135 logging.getLogger(__name__).info( 

136 "Simulation skipped (enable_simulation=False). " 

137 "Pipeline bubble and step time estimates are not available." 

138 ) 

139 

140 return result 

141 

142 @staticmethod 

143 def _run_simulation( 

144 result: PPStrategyResult, 

145 pipeline: Any, 

146 yaml_config: Any, 

147 constant_memory: int, 

148 sim_comm_time: float, 

149 ) -> None: 

150 """Run pipeline simulator and merge results into *result* in-place. 

151 

152 Creates a :class:`PPSimulator`, runs the post-ILP simulation, 

153 and merges the simulator output fields into the provided 

154 ``result``. On failure, marks ``simulation_status`` as 

155 ``"failed"`` with a descriptive error. 

156 

157 Args: 

158 result: Strategy result from :meth:`PPBalancer.balance_with_ilp` 

159 (modified in-place with simulation metrics). 

160 pipeline: Solved :class:`SappPipeline` instance. 

161 yaml_config: YAML configuration with pipeline topology. 

162 constant_memory: Constant memory per stage (MB). 

163 sim_comm_time: P2P communication time between adjacent stages 

164 (ms) for the simulator. 

165 """ 

166 pp_sim = PPSimulator( 

167 pipeline=pipeline, 

168 yaml_config=yaml_config, 

169 constant_memory=constant_memory, 

170 ) 

171 

172 sim_result = pp_sim.simulate_from_ilp(sim_comm_time=sim_comm_time) 

173 

174 if sim_result is None: 

175 result.simulation_status = "failed" 

176 result.simulation_error = ( 

177 "ILP simulation returned None (e.g. micro_batch_num < pp_degree)" 

178 ) 

179 return 

180 

181 real_bubble_val = sim_result.simulator_bubbles.get("real", 0.0) 

182 if not math.isfinite(sim_result.simulator_end_time) or not math.isfinite(real_bubble_val): 

183 result.simulation_status = "failed" 

184 result.simulation_error = ( 

185 "Simulation produced non-finite results (e.g. total compute time is zero)" 

186 ) 

187 return 

188 

189 result.simulation_status = sim_result.simulation_status 

190 result.simulation_error = sim_result.simulation_error 

191 result.simulator_end_time = sim_result.simulator_end_time 

192 result.simulator_bubbles = sim_result.simulator_bubbles 

193 result.simulator_peak_memory = sim_result.simulator_peak_memory 

194 

195 pipeline_bubble = sim_result.simulator_bubbles.get("real") 

196 if pipeline_bubble is None: 

197 logging.getLogger(__name__).warning( 

198 "Simulator bubbles dict missing 'real' key. " 

199 "Available keys: %s", 

200 list(sim_result.simulator_bubbles.keys()), 

201 ) 

202 result.pipeline_bubble = pipeline_bubble