Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_ppb / pp_modeling / pp_structs.py: 100%

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"""Core data structures and types for Pipeline Parallelism modeling.""" 

16 

17from __future__ import annotations 

18 

19from dataclasses import dataclass, field 

20from typing import Dict, List, Optional, Any, Tuple 

21 

22from hyper_parallel.auto_parallel.sapp_ppb.utils.recompute import TYPE as RecomputeType 

23 

24 

25@dataclass 

26class PPBOutput: 

27 """Output from PPB load balancing algorithm. 

28 

29 :class:`PPStrategyResult` inherits this class and adds 

30 ``pp_degree``, ``micro_batch_num``, ``vpp_less_memory``, and 

31 ``pipeline_bubble``. 

32 

33 Args: 

34 stage_partition: Stage partition result. When 

35 ``num_of_interleave == 1``, each entry is a physical stage 

36 list of ``(layer_id, RecomputeType)`` tuples. When 

37 ``num_of_interleave > 1``, the list has ``vpp * pp`` 

38 entries representing **virtual stages**: entry 

39 ``v * pp + s`` corresponds to VPP chunk ``v`` on physical 

40 stage ``s``. 

41 layer_offset: Per-group layer offset. Key is the BODY group name 

42 from the sapp-ppb ``Layer.name_`` attribute (see 

43 :attr:`PPStrategyResult.layer_offset` for naming rules); value is 

44 the offset matrix with shape ``[vpp][pp]``. 

45 is_feasible: Whether the ILP solution is feasible. This field 

46 reflects **only** the ILP solver outcome; it is **not** 

47 affected by pipeline simulator success or failure. 

48 Defaults to ``False`` (fail-safe) — callers must 

49 explicitly set it to ``True`` when the solver finds a 

50 feasible solution. 

51 infeasibility_details: Details about infeasibility. 

52 is_successful: Whether the ILP solver found a successful solution. 

53 ``True`` when the solver proved optimality or found a feasible 

54 incumbent; ``False`` otherwise. Propagated to 

55 :attr:`PPStrategyResult.is_successful`. 

56 simulation_status: Pipeline simulator execution status. One of 

57 ``"not_run"`` (simulator was not invoked), ``"success"`` 

58 (simulator completed), or ``"failed"`` (simulator could not 

59 run, e.g. ``micro_batch_num < pp_degree``). Independent of 

60 ``is_feasible``. 

61 simulation_error: Human-readable error message when 

62 ``simulation_status`` is ``"failed"``. 

63 simulator_end_time: Pipeline step time (ms) from the 

64 simulator. 0.0 when the simulator was not run or failed. 

65 simulator_bubbles: Per-type bubble ratios from the simulator. 

66 Empty dict when the simulator was not run. 

67 simulator_peak_memory: Per-stage peak memory (MB) from the 

68 simulator. Empty list when the simulator was not run. 

69 num_of_interleave: VPP interleaving factor. When > 1, 

70 ``stage_partition`` contains ``vpp * pp_degree`` entries 

71 (virtual stages) rather than ``pp_degree`` entries. 

72 

73 Example: 

74 >>> output = PPBOutput( 

75 ... stage_partition=[[(0, RecomputeType.NONE), (1, RecomputeType.NONE)], 

76 ... [(2, RecomputeType.SLCT), (3, RecomputeType.NONE)]], 

77 ... is_successful=True, 

78 ... ) 

79 """ 

80 

81 stage_partition: List[List[Tuple[int, RecomputeType]]] = field(default_factory=list) 

82 layer_offset: Dict[str, List[List[int]]] = field(default_factory=dict) 

83 is_feasible: bool = False 

84 infeasibility_details: Dict[str, Any] = field(default_factory=dict) 

85 is_successful: bool = False 

86 simulator_end_time: float = 0.0 

87 simulator_bubbles: Dict[str, float] = field(default_factory=dict) 

88 simulator_peak_memory: List[float] = field(default_factory=list) 

89 simulation_status: str = "not_run" 

90 simulation_error: Optional[str] = None 

91 num_of_interleave: int = 1 

92 

93 

94@dataclass 

95class PPStrategyResult(PPBOutput): 

96 """Pipeline parallelism strategy evaluation result. 

97 

98 Inherits :class:`PPBOutput` and adds pipeline-topology fields and 

99 the convenience ``pipeline_bubble`` accessor. 

100 

101 Inherited fields (from :class:`PPBOutput`): 

102 

103 * ``stage_partition`` — per-stage ``(layer_id, RecomputeType)`` tuples 

104 * ``layer_offset`` — per-group offset matrix ``[vpp][pp]`` 

105 * ``is_feasible`` — ILP feasibility flag 

106 * ``infeasibility_details`` — reason / solver status when infeasible 

107 * ``is_successful`` — ``True`` when ILP found a usable solution 

108 * ``simulator_end_time`` — estimated step time in ms (0.0 when not 

109 run or failed; replaces the former ``estimated_step_time`` field) 

110 * ``simulator_bubbles`` — per-type bubble ratios 

111 * ``simulator_peak_memory`` — per-stage peak memory (MB) 

112 * ``simulation_status`` — ``"not_run"`` / ``"success"`` / ``"failed"`` 

113 * ``simulation_error`` — error message when simulation fails 

114 * ``num_of_interleave`` — VPP interleaving factor 

115 

116 Args: 

117 pp_degree: Number of pipeline stages. 

118 micro_batch_num: Number of micro batches. 

119 vpp_less_memory: Whether the less-memory VPP schedule (``vpp2``) 

120 was used during optimization. Downstream consumers need 

121 this to reconstruct the correct pipeline schedule. 

122 pipeline_bubble: Pipeline bubble ratio (0.0 to 1.0). This is a 

123 **ratio**, not an absolute time in ms. ``None`` when the 

124 pipeline simulator did not produce a usable result (e.g. 

125 ``simulation_status`` is not ``"success"``). 

126 

127 Example: 

128 >>> result = PPStrategyResult( 

129 ... pp_degree=2, 

130 ... micro_batch_num=4, 

131 ... pipeline_bubble=0.25, 

132 ... is_successful=True, 

133 ... ) 

134 """ 

135 

136 pp_degree: int = 0 

137 micro_batch_num: int = 1 

138 vpp_less_memory: bool = False 

139 pipeline_bubble: Optional[float] = None