Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / common / layer_order.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"""Where the layers of a pipeline partition sit, in model order. 

16 

17A partition lists each stage's chunks and each chunk's layers. Model order 

18is the order of the layers in the model, which a layer's kind and its 

19recompute range count in: chunk by chunk across the stages, and back up 

20them in the second chunk of a V schedule. 

21""" 

22from typing import Any, List, Optional, Sequence, Tuple 

23 

24from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType 

25 

26# The layer type the partition places for each recompute option. 

27_TYPES = { 

28 "none": LayerType.NOT_REC_LAYER, 

29 "full": LayerType.FULL_REC_LAYER, 

30 "selective": LayerType.SEL_REC_LAYER, 

31} 

32 

33 

34def get_model_order(cfg: Any, stages: List) -> List[Tuple[int, int, int]]: 

35 """Positions ``(stage, chunk, index)`` of the regular layers, in model order. 

36 

37 Model order runs chunk by chunk across the stages, and back up them in 

38 the second chunk of a V schedule; walking the stages one at a time is 

39 only the same order without interleaving. This is the order the memory 

40 backbone gives layers their groups in. 

41 """ 

42 positions = [ 

43 (stage_id, chunk_id, lay_id) 

44 for chunk_id in range(cfg.vp) 

45 for stage_id in range(cfg.p) 

46 for lay_id, layer in enumerate(stages[stage_id][chunk_id]) 

47 if layer not in (LayerType.EMBEDDING_LAYER, LayerType.OUTPUT_LAYER) 

48 ] 

49 if cfg.pp_sched == "zero_bubble_v": 

50 # First-chunk entries less the embedding, as the memory backbone counts. 

51 first = sum(len(stage[0]) for stage in stages) - 1 

52 positions = positions[:first] + positions[first:][::-1] 

53 return positions 

54 

55 

56def stated_recompute(ccfg: Any) -> Optional[Tuple[Any, ...]]: 

57 """Return the recompute ranges a config states, or ``None`` when it states none. 

58 

59 An empty tuple is a statement too: no layer recomputes. 

60 """ 

61 ranges = getattr(ccfg, "recompute_ranges", None) 

62 return ranges if isinstance(ranges, tuple) else None 

63 

64 

65def layer_switches(ccfg: Any) -> Optional[Tuple[Any, ...]]: 

66 """Return each layer's own recompute switches in model order, or ``None`` where ``rec_op`` holds the one setting.""" 

67 switches = getattr(ccfg, "layer_switches", None) 

68 return switches if isinstance(switches, tuple) else None 

69 

70 

71def layer_recompute_types(ranges: Sequence[Any], layers: int) -> List[LayerType]: 

72 """Return the layer type of each of *layers* layers in model order. 

73 

74 Each layer takes the option of the recompute range that covers it, and 

75 a layer no range covers is not recomputed. 

76 """ 

77 types = [LayerType.NOT_REC_LAYER] * layers 

78 for item in ranges: 

79 stop = layers if item.count is None else min(layers, item.first + item.count) 

80 for index in range(item.first, stop): 

81 types[index] = _TYPES[item.option] 

82 return types