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

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 simulator — pipeline parallelism schedule simulation.""" 

16 

17from __future__ import annotations 

18 

19import logging 

20from typing import Any, Optional 

21 

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

23 SAPP_PPB_AVAILABLE, 

24) 

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

26 YamlOptimizationConfig, 

27) 

28from hyper_parallel.auto_parallel.sapp_ppb.pp_modeling.pp_structs import PPBOutput 

29 

30logger = logging.getLogger(__name__) 

31 

32 

33class PPSimulator: 

34 """Pipeline parallelism schedule simulator. 

35 

36 Wraps :class:`SappPipeline.simulate` and provides high-level methods 

37 for simulating ILP-optimized pipeline schedules. 

38 

39 Args: 

40 pipeline: A solved :class:`SappPipeline` instance. 

41 yaml_config: YAML configuration with pipeline topology. 

42 constant_memory: Constant memory per stage (MB). 

43 

44 Example: 

45 >>> sim = PPSimulator(pipeline, yaml_config, constant_mem) 

46 >>> result = sim.simulate_from_ilp(sim_comm_time=0.1) 

47 """ 

48 

49 def __init__( 

50 self, 

51 pipeline: Any, 

52 yaml_config: YamlOptimizationConfig, 

53 constant_memory: int, 

54 ) -> None: 

55 """Initialize PPSimulator. 

56 

57 Args: 

58 pipeline: A solved :class:`SappPipeline` instance. 

59 yaml_config: YAML configuration with pipeline topology. 

60 constant_memory: Constant memory per stage (MB). 

61 """ 

62 self._pipeline = pipeline 

63 self._yaml_config = yaml_config 

64 self._constant_memory = constant_memory 

65 

66 def simulate_from_ilp( 

67 self, 

68 sim_comm_time: float = 0.0, 

69 ) -> Optional[PPBOutput]: 

70 """Run PipelineSimulator using ILP solver output for accurate step time. 

71 

72 Delegates to :meth:`SappPipeline.simulate` which internally calls 

73 ``get_fw_time()``, ``get_recompute_time()``, 

74 ``get_memory_activation()``, and ``get_memory_parameter()`` 

75 to build and run the :class:`PipelineSimulator`. The resulting 

76 ``end_time`` reflects the true pipeline schedule (1F1B / VPP 

77 with warmup-steady-cooldown phases, P2P communication, bubble 

78 overlap) — far more accurate than 

79 ``max_stage_time × micro_batch_num``. 

80 

81 Args: 

82 sim_comm_time: P2P communication time between adjacent stages 

83 (ms). Passed to the pipeline simulator only; does NOT 

84 affect ILP optimization. Default 0.0 (no communication 

85 delay). 

86 

87 Returns: 

88 :class:`PPBOutput` with ``simulation_status="success"`` and 

89 populated ``simulator_end_time``, ``simulator_bubbles``, 

90 ``simulator_peak_memory``; or ``None`` if the pipeline has 

91 not been solved yet or the simulation cannot run. 

92 

93 Example: 

94 >>> sim = PPSimulator(pipeline, yaml_config, constant_mem) 

95 >>> result = sim.simulate_from_ilp(sim_comm_time=0.1) 

96 >>> result.simulator_end_time 

97 1234.5 

98 """ 

99 if not SAPP_PPB_AVAILABLE or self._pipeline is None: 

100 return None 

101 

102 if self._yaml_config.micro_batch_num < self._yaml_config.pp_degree: 

103 logger.warning( 

104 "micro_batch_num (%d) < pp_degree (%d); simulator skipped.", 

105 self._yaml_config.micro_batch_num, self._yaml_config.pp_degree, 

106 ) 

107 return None 

108 

109 try: 

110 end_time = self._pipeline.simulate( 

111 show=False, comm_time=sim_comm_time, 

112 ) 

113 

114 if end_time is None or end_time <= 0: 

115 return None 

116 

117 sim_instance = self._pipeline.simulator 

118 return PPBOutput( 

119 simulation_status="success", 

120 simulator_end_time=sim_instance.end_time, 

121 simulator_bubbles=dict(sim_instance.bubbles), 

122 simulator_peak_memory=list(sim_instance.peak_memory), 

123 ) 

124 except ValueError as exc: 

125 logger.warning("Simulator failed with ValueError: %s", exc) 

126 return None 

127 except Exception as exc: 

128 from hyper_parallel.auto_parallel.sapp_ppb.simulator.causal_error import CausalCommError, CausalError # pylint: disable=C0415 

129 if isinstance(exc, (CausalCommError, CausalError)): 

130 logger.warning("Simulator failed with %s: %s", type(exc).__name__, exc) 

131 return None 

132 raise