Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_ppb / run_pipeline_balance.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"""CLI entrypoint for SAPP-PPB pipeline balancing.""" 

16import argparse 

17import os 

18import sys 

19 

20from hyper_parallel.auto_parallel.sapp_ppb.sapp.sapp_pipeline import SappPipeline 

21from hyper_parallel.auto_parallel.sapp_ppb.utils import interactive 

22from hyper_parallel.auto_parallel.sapp_ppb.utils.compute_memory import compute_memories 

23from hyper_parallel.auto_parallel.sapp_ppb.utils.config import initialize_layer_json 

24from hyper_parallel.auto_parallel.sapp_ppb.utils.layer import generate_layers_list 

25from hyper_parallel.auto_parallel.sapp_ppb.utils.logger import logger 

26 

27 

28def _str2bool(value: str) -> bool: 

29 """Parse a truthy string value coming from ``argparse``.""" 

30 return str(value).lower() in ('true', '1', 'yes') 

31 

32 

33def build_arg_parser() -> argparse.ArgumentParser: 

34 """Build the argument parser for the pipeline-balance CLI. 

35 

36 Returns: 

37 Configured :class:`argparse.ArgumentParser` instance. 

38 """ 

39 parser = argparse.ArgumentParser( 

40 prog='SAPP AutoBalancing', 

41 description='Balance layers onto pipeline stages, considering recomputation and interleaving', 

42 epilog='') 

43 

44 # Pipeline info 

45 parser.add_argument('-s', '--stage', type=int, default=4, help="Number of stages") 

46 parser.add_argument('-mb', '--micro_batch', type=int, default=4, help="Number of micro batch") 

47 parser.add_argument('-i', '--interleave_degree', type=int, default=1, help="Interleave level") 

48 

49 # Memory size 

50 parser.add_argument('-mem', '--max_memory', type=int, default=56000, 

51 help="Maximum memory available (MB)") 

52 parser.add_argument('-lm', '--less_memory', type=_str2bool, default=False, 

53 help="Compute Memory with 'Less Memory interleave' option") 

54 parser.add_argument('-dual', '--dualpipe_v', type=_str2bool, default=False, 

55 help="Compute Memory with 'DualpipeV' option") 

56 parser.add_argument('-mc', '--constant_memory', type=int, default=0, 

57 help="Constant memory per stages") 

58 

59 parser.add_argument('-o', '--output_folder', type=str, default="./output", 

60 help="output files location") 

61 

62 # Model info 

63 parser.add_argument('-m', '--model_name', type=str, default="model_name", help="") 

64 

65 # Search time 

66 parser.add_argument('-t', '--time_limit', type=int, default=90, 

67 help="Limitation on searching time") 

68 

69 # Optimization level 

70 parser.add_argument('-O', '--optimization_level', type=int, default=1, 

71 help="Defines optimization level when Stage (S) = Micro Batch number (M). " 

72 "0 for same approach as M > S. " 

73 "1 (default) generally better. " 

74 "2 better for memory constrained cases.") 

75 

76 # Simulate naive or manual config 

77 parser.add_argument('-naive', '--simulate_naive', type=_str2bool, default=False, 

78 help="Simulate naive configs") 

79 parser.add_argument('-manual', '--manual_config', type=str, default=None, 

80 help="Path of manual config") 

81 

82 # Layer info 

83 parser.add_argument('-lf', '--layer_folder', type=str, default="./layers/", 

84 help="Path to the layer folder") 

85 parser.add_argument('-dump', '--dump_layer', type=_str2bool, default=False, 

86 help="Dump the layers") 

87 

88 # For Computation of memory 

89 parser.add_argument('-mf', '--memory_folder', type=str, default="./memory/", 

90 help="Path to the profiler memory folder") 

91 

92 # For Initialization 

93 parser.add_argument('-init', '--init', type=str, default=None, 

94 help="Path to the init file") 

95 

96 # Computation argument 

97 parser.add_argument('-cm', '--compute_memory', type=_str2bool, default=False, 

98 help="Parse Mindspore log to generate MEMORY of the layer (unavailable)") 

99 parser.add_argument('-exec', '--exec', type=_str2bool, default=True, 

100 help="Compute solver") 

101 return parser 

102 

103 

104def _resolve_path(base_dir: str, path: str) -> str: 

105 """Return ``path`` resolved relative to ``base_dir`` unless it is already absolute.""" 

106 if os.path.isabs(path): 

107 return path 

108 return os.path.join(base_dir, path) 

109 

110 

111def run(args: argparse.Namespace, base_dir: str) -> None: 

112 """Execute the pipeline balancing workflow for the given arguments. 

113 

114 Args: 

115 args (argparse.Namespace): Parsed CLI arguments. 

116 base_dir (str): Directory used to resolve relative input / output paths. 

117 """ 

118 if args.init: 

119 init_file = _resolve_path(base_dir, args.init) 

120 initialize_layer_json(args.model_name, init_file) 

121 

122 output_folder = _resolve_path(base_dir, args.output_folder) 

123 os.makedirs(output_folder, exist_ok=True) 

124 

125 manual_config = None 

126 if args.manual_config: 

127 candidate = _resolve_path(base_dir, args.manual_config) 

128 if candidate.endswith(('yaml', 'yml')): 

129 manual_config = candidate 

130 

131 layers = generate_layers_list(args.layer_folder, args.model_name) 

132 if args.compute_memory: 

133 layers = compute_memories(layers=layers, memory_folder=args.memory_folder, 

134 number_of_stage=args.stage) 

135 for layer in layers: 

136 logger.output("%s", layer) 

137 

138 if args.dump_layer: 

139 for layer in layers: 

140 layer.dump() 

141 

142 pipe = SappPipeline(model_name=args.model_name, num_of_stage=args.stage, 

143 num_of_micro_batch=args.micro_batch, max_memory=args.max_memory, 

144 layers=layers, num_of_interleave=args.interleave_degree, 

145 vpp_less_memory=args.less_memory, dual=args.dualpipe_v, 

146 constant_memory=args.constant_memory, 

147 optimization_level=args.optimization_level) 

148 

149 pipe.construct_problem(solver="pulp") 

150 

151 if args.exec: 

152 pipe.solve_problem(time_limit=args.time_limit, dump_folder=output_folder) 

153 pipe.print_yaml_results() 

154 total_time = pipe.simulate(show=True, file_name=os.path.join(output_folder, "result.svg")) 

155 

156 logger.output("total_time: %d", total_time) 

157 logger.output("time: %s", pipe.get_time()) 

158 logger.output("mem_par: %s", pipe.get_memory_parameter()) 

159 logger.output("mem_act: %s", pipe.get_memory_activation()) 

160 

161 if manual_config: 

162 logger.output("Simulating manual configs") 

163 pipe.simulate_comparison(manual_config, output_folder) 

164 if args.simulate_naive: 

165 logger.output("Simulating naive configs") 

166 pipe.simulate_naive(layers, output_folder) 

167 elif manual_config: 

168 logger.output("Simulating manual configs") 

169 pipe.simulate_only_manual(manual_config, output_folder) 

170 

171 

172def main() -> None: 

173 """Entry point invoked when the module is run as a script.""" 

174 if len(sys.argv) == 1: 

175 interactive.main() 

176 return 

177 parser = build_arg_parser() 

178 args = parser.parse_args() 

179 run(args, base_dir=os.path.dirname(os.path.abspath(__file__))) 

180 

181 

182if __name__ == "__main__": 

183 main()