Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / compile.py: 25%

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"""Compile Transformer decoder layers as independent graph segments.""" 

16 

17import logging 

18from collections.abc import Iterable, Mapping 

19from typing import Any, Optional, Union 

20 

21import torch 

22from torch import nn 

23 

24from hyper_parallel.models.build_options import CompileConfig 

25from hyper_parallel.distributed._builder.fsdp_adapter import FSDP2Manager 

26 

27logger = logging.getLogger(__name__) 

28 

29 

30_MAPPING_GET_POLYFILL_INSTALLED = False 

31 

32# No per-model layer-path table: the model-owned ``get_compile_layers()`` 

33# contract comes first, and the generic HF-convention fallback paths below 

34# cover every supported family. A family whose container lives elsewhere 

35# declares the contract on its model class instead of registering here. 

36_GENERIC_LAYER_PATHS = ("layers", "model.layers") 

37 

38 

39def _get_attribute(root: Any, path: str) -> Any: 

40 """Read a dotted module path, including numeric ModuleList indexes.""" 

41 value = root 

42 for part in path.split("."): 

43 value = value[int(part)] if part.isdigit() else getattr(value, part) 

44 return value 

45 

46 

47def _is_named_layer(item: Any) -> bool: 

48 """Return whether an item is a valid ``(fqn, module)`` pair.""" 

49 if not isinstance(item, tuple) or len(item) != 2: 

50 return False 

51 name, layer = item 

52 return isinstance(name, str) and isinstance(layer, nn.Module) 

53 

54 

55def _normalize_declared_layers(declared: Any) -> list[tuple[str, nn.Module]]: 

56 """Normalize the model-owned decoder-layer contract.""" 

57 if isinstance(declared, nn.ModuleList): 

58 return [(str(index), layer) for index, layer in enumerate(declared)] 

59 if not isinstance(declared, Iterable) or isinstance(declared, (str, bytes)): 

60 raise TypeError( 

61 "get_compile_layers() must return an iterable of modules or " 

62 "(fqn, module) pairs" 

63 ) 

64 

65 layers = [] 

66 for index, item in enumerate(declared): 

67 if isinstance(item, nn.Module): 

68 layers.append((str(index), item)) 

69 continue 

70 if _is_named_layer(item): 

71 layers.append(item) 

72 continue 

73 raise TypeError( 

74 "get_compile_layers() entries must be modules or (fqn, module) pairs, " 

75 f"but entry {index} is {type(item).__name__}" 

76 ) 

77 return layers 

78 

79 

80def _qualify_declared_layers( 

81 model: nn.Module, 

82 layers: list[tuple[str, nn.Module]], 

83) -> list[tuple[str, nn.Module]]: 

84 """Replace generated numeric names with the modules' model-owned FQNs.""" 

85 module_names = {id(module): name for name, module in model.named_modules()} 

86 return [ 

87 (module_names.get(id(layer), name) if name.isdigit() else name, layer) 

88 for name, layer in layers 

89 ] 

90 

91 

92def _layers_from_path(model: nn.Module, path: str) -> list[tuple[str, nn.Module]]: 

93 """Return indexed layers from one declared container path.""" 

94 try: 

95 container = _get_attribute(model, path) 

96 except (AttributeError, IndexError, KeyError, TypeError): 

97 return [] 

98 if not isinstance(container, (nn.ModuleList, nn.Sequential)): 

99 return [] 

100 return [(f"{path}.{index}", layer) for index, layer in enumerate(container)] 

101 

102 

103def get_compile_layers(model: nn.Module) -> list[tuple[str, nn.Module]]: 

104 """Return stable decoder-layer segments declared by a supported model. 

105 

106 Resolution order: the model-owned ``get_compile_layers()`` contract 

107 first, then the generic HF-convention container paths 

108 (``_GENERIC_LAYER_PATHS``). A family whose container follows neither 

109 declares ``get_compile_layers()`` on its model class. 

110 """ 

111 declared_getter = getattr(model, "get_compile_layers", None) 

112 if callable(declared_getter): 

113 layers = _qualify_declared_layers( 

114 model, 

115 _normalize_declared_layers(declared_getter()), 

116 ) 

117 if layers: 

118 return layers 

119 

120 seen_containers = set() 

121 for path in _GENERIC_LAYER_PATHS: 

122 layers = _layers_from_path(model, path) 

123 if not layers: 

124 continue 

125 container_id = id(_get_attribute(model, path)) 

126 if container_id in seen_containers: 

127 continue 

128 seen_containers.add(container_id) 

129 return layers 

130 

131 raise ValueError( 

132 "compile is enabled, but the model exposes no decoder-layer compile contract; " 

133 "define get_compile_layers() or use a supported model layer container" 

134 ) 

135 

136 

137def _install_dynamo_mapping_get_polyfill() -> None: 

138 """Make ``Mapping.get`` traceable without changing Transformers source. 

139 

140 Transformers' attention registry calls ``Mapping.get`` from every decoder 

141 layer. TorchDynamo treats the standard-library implementation as a 

142 skipfile, so the call creates a graph break. ``substitute_in_graph`` only 

143 replaces the implementation while Dynamo inlines it; eager execution 

144 continues to use the original method. 

145 """ 

146 # Module-level idempotency flag: the polyfill must be installed at most 

147 # once per process, so a global statement is required here. 

148 global _MAPPING_GET_POLYFILL_INSTALLED # pylint: disable=global-statement 

149 if _MAPPING_GET_POLYFILL_INSTALLED: 

150 return 

151 

152 substitute_in_graph = getattr(torch.compiler, "substitute_in_graph", None) 

153 if substitute_in_graph is None: 

154 logger.warning( 

155 "Torch does not provide compiler.substitute_in_graph; " 

156 "Mapping.get graph breaks cannot be removed" 

157 ) 

158 _MAPPING_GET_POLYFILL_INSTALLED = True 

159 return 

160 

161 def _mapping_get(self, key, default=None): 

162 try: 

163 return self[key] 

164 except KeyError: 

165 return default 

166 

167 try: 

168 substitute_in_graph(Mapping.get)(_mapping_get) 

169 except ValueError as exc: 

170 if "Duplicate dispatch rule" not in str(exc): 

171 raise 

172 logger.debug("Mapping.get already has a TorchDynamo substitution") 

173 _MAPPING_GET_POLYFILL_INSTALLED = True 

174 

175 

176def resolve_compile_kwargs(config: CompileConfig) -> dict[str, Any]: 

177 """Convert ``CompileConfig`` into keyword arguments for ``Module.compile``.""" 

178 kwargs: dict[str, Any] = { 

179 "fullgraph": config.fullgraph, 

180 "dynamic": config.dynamic, 

181 } 

182 if config.backend is not None: 

183 kwargs["backend"] = config.backend 

184 if config.options: 

185 kwargs["options"] = dict(config.options) 

186 else: 

187 kwargs["mode"] = config.mode 

188 return kwargs 

189 

190 

191def apply_compile(model: nn.Module, config: CompileConfig) -> nn.Module: 

192 """Compile each decoder layer in place while preserving module identity.""" 

193 if not config.enabled: 

194 return model 

195 

196 _install_dynamo_mapping_get_polyfill() 

197 torch._dynamo.config.cache_size_limit = config.dynamo_cache_size_limit # pylint: disable=W0212 

198 compile_kwargs = resolve_compile_kwargs(config) 

199 layers = get_compile_layers(model) 

200 for layer_fqn, layer in layers: 

201 try: 

202 layer.compile(**compile_kwargs) 

203 except Exception as exc: 

204 raise RuntimeError(f"failed to compile segment {layer_fqn}") from exc 

205 

206 logger.info( 

207 "Compiled %d decoder layers (backend=%s, mode=%s, dynamic=%s, fullgraph=%s)", 

208 len(layers), 

209 config.backend, 

210 config.mode, 

211 config.dynamic, 

212 config.fullgraph, 

213 ) 

214 return model 

215 

216 

217__all__ = [ 

218 "apply_compile", 

219 "get_compile_layers", 

220 "resolve_compile_kwargs", 

221] 

222 

223 

224def _resolve_compile_config( 

225 compile_config: Optional[Union[CompileConfig, dict]], 

226 validate_placement: bool, 

227 fsdp2_manager: Optional[FSDP2Manager], 

228) -> tuple[Optional[CompileConfig], bool]: 

229 """Normalize compile configuration and validate its FSDP interaction. 

230 

231 Moved from ``_transformers/infrastructure.py`` (05 §15.2.1): compile 

232 option normalization and the fullgraph/FSDP conflict check follow the 

233 compile feature. 

234 """ 

235 if isinstance(compile_config, dict): 

236 compile_config = CompileConfig(enabled=True, **compile_config) 

237 if compile_config is not None and not isinstance(compile_config, CompileConfig): 

238 raise TypeError("compile_config must be a CompileConfig, mapping, or None") 

239 graph_mode_requested = bool( 

240 compile_config is not None and compile_config.selects_graph_compiler() 

241 ) 

242 compile_for_execution = bool( 

243 not validate_placement 

244 and compile_config is not None 

245 and compile_config.enabled 

246 and not graph_mode_requested 

247 ) 

248 if graph_mode_requested: 

249 logger.info( 

250 "compile.use_joint_graph=True selects graph-mode trainer; " 

251 "skipping decoder-layer compile", 

252 ) 

253 if validate_placement and compile_config is not None and compile_config.enabled: 

254 logger.info("Skipping decoder-layer compile during placement validation") 

255 if compile_for_execution and compile_config.fullgraph and isinstance(fsdp2_manager, FSDP2Manager): 

256 raise ValueError( 

257 "compile.fullgraph=True is incompatible with FSDP hooks kept eager " 

258 "by _dynamo_disable; set compile.fullgraph=False" 

259 ) 

260 return compile_config, compile_for_execution