Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / recompute_dimension.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"""Recompute as a dimension of the search: the activation checkpoint modes a candidate may run. 

16 

17A parallel degree is stated as one value, a list of candidates or ``auto``, 

18and the recompute dimension is stated the same way: ``full``, ``[off, full]`` 

19or ``auto``. Its values are HyperParallel's activation checkpoint modes. The 

20search prices every candidate under each mode the dimension allows and ranks 

21the pairs together, each with its own memory check, so a mode left out of the 

22list is never proposed: selective, say, which the estimate prices as a policy 

23that runs well would run, while HyperParallel's own selective path runs 

24slower than full recompute. 

25 

26The degrees are integers whose product is the device count, which the search 

27enumerates together; a mode is no factor of it, so the search runs once per 

28mode instead (:meth:`ParallelizeLayer.run_generation_to_ordering`). 

29""" 

30import copy 

31from typing import Any, List, Optional, Tuple 

32 

33from hyper_parallel.auto_parallel.sapp_nd.nd.common.config import Config 

34from hyper_parallel.auto_parallel.sapp_nd.nd.common.derive import hyper_rec_op 

35 

36RECOMPUTE_MODES = ("off", "selective", "full") 

37# The parsers whose recompute is HyperParallel's activation checkpoint mode, 

38# and whose selective switches follow from it alone. 

39HYPER_FRAMEWORKS = ("hyper_v2", "hyperparallel2") 

40# Older search configs state no recompute as "none". 

41_ALIASES = {"none": "off"} 

42 

43 

44def read_recompute_modes(stated: Any, where: str) -> Optional[Tuple[str, ...]]: 

45 """Read a recompute dimension stated the way a parallel degree is. 

46 

47 Args: 

48 stated: None where nothing states it; ``"auto"`` for every mode; one 

49 mode; or a list of modes. YAML reads an unquoted ``off`` as False, 

50 which counts as ``"off"``. 

51 where: Where it was stated, for the error message. 

52 

53 Returns: 

54 The modes in the order stated, each once, or None where nothing 

55 states the dimension and the search keeps the config's own recompute. 

56 

57 Raises: 

58 ValueError: An empty list, or a value that is neither a mode nor auto. 

59 """ 

60 if stated is None: 

61 return None 

62 listed = list(stated) if isinstance(stated, (list, tuple)) else [stated] 

63 if [str(value).strip().lower() for value in listed] == ["auto"]: 

64 return RECOMPUTE_MODES 

65 modes: List[str] = [] 

66 for value in listed: 

67 mode = "off" if value is False else str(value).strip().lower() 

68 mode = _ALIASES.get(mode, mode) 

69 if mode not in RECOMPUTE_MODES: 

70 raise ValueError( 

71 f"{where}: {value!r} is not a recompute mode; expected one of " 

72 f"{', '.join(RECOMPUTE_MODES)}, a list of them, or auto" 

73 ) 

74 if mode not in modes: 

75 modes.append(mode) 

76 if not modes: 

77 raise ValueError(f"{where}: an empty list allows no recompute mode") 

78 return tuple(modes) 

79 

80 

81def read_activation_checkpoint_mode(stated: Any, where: str) -> str: 

82 """Read one activation checkpoint mode the way the trainer reads ``activation_checkpoint.mode``. 

83 

84 Nothing stated is the trainer's default, ``"off"``, as is the older 

85 ``"none"``. YAML reads an unquoted ``off`` as False, which the trainer's 

86 config resolver turns back into ``"off"``, as here. Anything else that 

87 is not one of :data:`RECOMPUTE_MODES` is refused, True included: ND used 

88 to price a mode it did not know as no recompute, so a typo of ``full`` 

89 was costed at the off memory, and a False fell through to the legacy 

90 default. 

91 

92 Args: 

93 stated: The mode as the yaml states it. 

94 where: Where it was stated, for the error message. 

95 

96 Returns: 

97 One of :data:`RECOMPUTE_MODES`. 

98 

99 Raises: 

100 ValueError: A list, auto, True, or a name that is no mode. 

101 """ 

102 if stated is None: 

103 return "off" 

104 if isinstance(stated, (list, tuple)) or str(stated).strip().lower() == "auto": 

105 raise ValueError(f"{where}: {stated!r} is not one activation checkpoint mode; expected one of " 

106 f"{', '.join(RECOMPUTE_MODES)}") 

107 return read_recompute_modes(stated, where)[0] 

108 

109 

110def _mode_configs(ccfg: Any) -> List[Any]: 

111 """The configs a mode is stated on: the model's, and each submodule's of a multimodal one.""" 

112 configs = [ccfg] 

113 if getattr(ccfg, "multimodal", False): 

114 configs += [ccfg.mm_ccfgs[name] for name in ccfg.mm_order] 

115 return configs 

116 

117 

118def state_recompute_mode(ccfg: Any, mode: str) -> None: 

119 """State one activation checkpoint mode on a model's cost model config, as the HyperParallel parsers read it. 

120 

121 Each candidate's full recompute comes from the pipeline balancing, which 

122 the search states separately (:meth:`GlobalConfig.state_recompute`). 

123 

124 Args: 

125 ccfg: The whole model's cost model config. 

126 mode: One of :data:`RECOMPUTE_MODES`. 

127 """ 

128 for config in _mode_configs(ccfg): 

129 config.full_rec = mode == "full" 

130 config.sel_rec = mode == "selective" 

131 config.rec_op = Config(hyper_rec_op(config.sel_rec)) 

132 

133 

134def parsed_recompute(ccfg: Any) -> List[Tuple[Any, Any, Any, Any]]: 

135 """Each config's recompute as its parser stated it, for :func:`restore_recompute`.""" 

136 return [(config, copy.deepcopy(config.full_rec), copy.deepcopy(config.sel_rec), 

137 copy.deepcopy(getattr(config, "rec_op", None))) 

138 for config in _mode_configs(ccfg)] 

139 

140 

141def restore_recompute(parsed: List[Tuple[Any, Any, Any, Any]]) -> None: 

142 """Give each config back the recompute :func:`parsed_recompute` took from it.""" 

143 for config, full_rec, sel_rec, rec_op in parsed: 

144 config.full_rec = copy.deepcopy(full_rec) 

145 config.sel_rec = copy.deepcopy(sel_rec) 

146 config.rec_op = copy.deepcopy(rec_op)