Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / model_build_overlay.py: 98%

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"""Codegen-owned model observation layered over the native AutoModel API.""" 

16# pylint: disable=forbidden-backend-import 

17 

18from collections.abc import Sequence 

19from contextlib import contextmanager 

20from dataclasses import dataclass 

21import inspect 

22import random 

23from typing import Any, Iterator, Optional 

24 

25import torch 

26 

27from hyper_parallel.codegen.manager import ensure_codegen_artifact_from_snapshot 

28from hyper_parallel.codegen.session import CodegenBuildSession 

29from hyper_parallel.core.dtensor.init_weights import init_empty_weights 

30from hyper_parallel.distributed._builder.applier import ( 

31 _get_active_mesh, 

32 _preflight_compute_injection, 

33) 

34from hyper_parallel.models._transformers import auto_model, model_builder 

35from hyper_parallel.models import module_build, replacement as replacement_api 

36 

37 

38def _prepare_model_source_compat(adapter: Any) -> None: 

39 """Run an optional model-owned compatibility hook before construction.""" 

40 loader = getattr(adapter, "codegen", None) if adapter is not None else None 

41 provider = loader() if loader is not None else None 

42 prepare = getattr(provider, "prepare_model_import", None) 

43 if prepare is not None: 

44 prepare() 

45 

46 

47@dataclass 

48class CodegenModelBuildResult: 

49 """Result populated by one scoped offline model observation.""" 

50 

51 artifact_layout: Optional[Any] = None 

52 snapshot: Optional[Any] = None 

53 

54 

55@contextmanager 

56def _codegen_model_init_context() -> Iterator[None]: 

57 """Build codegen's observation model with parameters and buffers on meta.""" 

58 with init_empty_weights(include_buffers=True), auto_model._model_init_context( # pylint: disable=protected-access 

59 True 

60 ): 

61 yield 

62 

63 

64def _record_source( 

65 session: CodegenBuildSession, 

66 model: Any, 

67 hf_config: Any, 

68 pretrained_model_name_or_path: Any, 

69 is_hf_model: bool, 

70) -> None: 

71 """Record deterministic source identity from the constructed meta model.""" 

72 model_type = type(model) 

73 source_path = inspect.getsourcefile(model_type) 

74 architectures = list(getattr(hf_config, "architectures", None) or ()) 

75 architecture = architectures[0] if architectures else model_type.__name__ 

76 session.record_source({ 

77 "model_type": getattr(hf_config, "model_type", None), 

78 "architecture": architecture, 

79 "architectures": architectures, 

80 "model_class": f"{model_type.__module__}:{model_type.__qualname__}", 

81 "module_name": model_type.__module__, 

82 "file_path": source_path, 

83 "pretrained_model_name_or_path": ( 

84 None if pretrained_model_name_or_path is None else str(pretrained_model_name_or_path) 

85 ), 

86 "is_hf_model": bool(is_hf_model), 

87 }) 

88 

89 

90def _observe_replacements( 

91 model: Any, 

92 distributed_setup: Any, 

93 low_precision_config: Any, 

94 load_base_model: bool, 

95 session: CodegenBuildSession, 

96) -> Any: 

97 """Apply native replacement declarations while recording their result.""" 

98 specs = list(getattr(distributed_setup, "module_replacements", None) or ()) 

99 original_compile = module_build.compile_module_replacements 

100 

101 def _observe_compile(*args: Any, **kwargs: Any) -> Any: 

102 plan = original_compile(*args, **kwargs) 

103 specs[:] = plan.specs 

104 return plan 

105 

106 module_build.compile_module_replacements = _observe_compile 

107 try: 

108 with _observe_replacement_installations(specs) as applications: 

109 model, _ = model_builder._maybe_apply_module_replacements( # pylint: disable=protected-access 

110 model, distributed_setup, low_precision_config, load_base_model, 

111 ) 

112 except replacement_api.ModuleReplacementMatchError as exc: 

113 session.record_replacements( 

114 declarations=replacement_api.describe_replacement_declarations(specs), 

115 applications=(), unmatched=exc.unmatched, 

116 ) 

117 raise 

118 finally: 

119 module_build.compile_module_replacements = original_compile 

120 session.record_replacements( 

121 declarations=replacement_api.describe_replacement_declarations(specs), 

122 applications=applications, unmatched=(), 

123 ) 

124 return model 

125 

126 

127@contextmanager 

128def _observe_replacement_installations(specs: Sequence[Any]) -> Iterator[list[dict[str, Any]]]: 

129 """Record metadata only after the shared installer commits a prepared candidate.""" 

130 applications: list[dict[str, Any]] = [] 

131 spec_indexes = {id(spec): index for index, spec in enumerate(specs)} 

132 pending = {} 

133 original_prepare = replacement_api._prepare_replacement_target # pylint: disable=protected-access 

134 original_install = replacement_api._install_module_candidates # pylint: disable=protected-access 

135 

136 def _observe_prepare(target: Any, *args: Any, **kwargs: Any) -> Any: 

137 candidate, transforms = original_prepare(target, *args, **kwargs) 

138 if id(target.spec) not in spec_indexes: 

139 spec_indexes.update({id(spec): index for index, spec in enumerate(specs)}) 

140 index = spec_indexes.setdefault(id(target.spec), len(spec_indexes)) 

141 pending[(target.module_fqns, id(candidate))] = { 

142 "rule_id": replacement_api._replacement_rule_id(target.spec, index), # pylint: disable=protected-access 

143 "module_fqns": list(target.module_fqns), 

144 "primary_fqn": target.module_fqns[0], 

145 "source_type": replacement_api._symbol_path(type(target.source)), 

146 "replacement_type": replacement_api._symbol_path(type(candidate)), 

147 "weight_transforms": [replacement_api._symbol_path(type(transform)) for transform in transforms], 

148 } 

149 return candidate, transforms 

150 

151 def _observe_install(root: Any, candidates: Any) -> Any: 

152 candidates = tuple(candidates) 

153 result = original_install(root, candidates) 

154 for fqns, _, candidate in candidates: 

155 record = pending.pop((fqns, id(candidate)), None) 

156 if record is not None: 

157 applications.append(record) 

158 return result 

159 

160 replacement_api._prepare_replacement_target = _observe_prepare # pylint: disable=protected-access 

161 replacement_api._install_module_candidates = _observe_install # pylint: disable=protected-access 

162 try: 

163 yield applications 

164 finally: 

165 replacement_api._prepare_replacement_target = original_prepare # pylint: disable=protected-access 

166 replacement_api._install_module_candidates = original_install # pylint: disable=protected-access 

167 

168 

169def _observe_plan( 

170 model: Any, 

171 mesh: Any, 

172 sharding_planner: Any, 

173 is_hf_model: bool, 

174 session: CodegenBuildSession, 

175) -> None: 

176 """Derive a native plan without applying distributed runtime operations.""" 

177 model_sharding_requested = ( 

178 mesh is not None 

179 and any(size > 1 for size in (mesh.tp_size, mesh.cp_size, mesh.ep_size)) 

180 ) 

181 if sharding_planner is None: 

182 session.mark_not_required("planner_unavailable") 

183 return 

184 if mesh is None: 

185 session.mark_not_required("mesh_unavailable") 

186 return 

187 if mesh.device_mesh is None: 

188 session.mark_not_required("device_mesh_unavailable") 

189 return 

190 if (is_hf_model and not model_sharding_requested 

191 and not model_builder._has_module_owned_layouts(model)): # pylint: disable=protected-access 

192 if session.has_replacement_applications: 

193 session.observe_plan(None, None, mesh) 

194 else: 

195 session.mark_not_required("pure_data_parallel") 

196 return 

197 plan = sharding_planner.plan( 

198 model, 

199 mesh.device_mesh, 

200 tp_size=mesh.tp_size, 

201 cp_size=mesh.cp_size, 

202 ep_size=mesh.ep_size, 

203 sequence_parallel=mesh.sequence_parallel, 

204 loss_parallel=mesh.loss_parallel, 

205 ) 

206 active_mesh = _get_active_mesh(mesh.device_mesh, plan.mesh_dim_names) 

207 _preflight_compute_injection(plan, active_mesh, model=model) 

208 session.observe_plan(plan, model, mesh) 

209 

210 

211def _publish_session( 

212 session: CodegenBuildSession, 

213 artifact_dir: Optional[str], 

214 result: CodegenModelBuildResult, 

215) -> Any: 

216 """Publish an offline snapshot while preserving host RNG state.""" 

217 cpu_rng_state = torch.get_rng_state() 

218 python_rng_state = random.getstate() 

219 

220 def _publish(snapshot: Any) -> None: 

221 """Publish one snapshot and restore RNG state afterward.""" 

222 try: 

223 result.artifact_layout = ensure_codegen_artifact_from_snapshot(snapshot, artifact_dir) 

224 finally: 

225 torch.set_rng_state(cpu_rng_state) 

226 random.setstate(python_rng_state) 

227 

228 return session.finish(_publish) 

229 

230 

231def _codegen_build_model(session: CodegenBuildSession, artifact_dir: Optional[str], result: CodegenModelBuildResult): 

232 """Create the classmethod used only while the CLI constructs a meta model.""" 

233 def _build_model( 

234 cls, 

235 pretrained_model_name_or_path, 

236 *model_args, 

237 is_hf_model, 

238 hf_config, 

239 mesh, 

240 sharding_planner, 

241 fsdp2_manager, 

242 peft_config, 

243 torch_dtype, 

244 attn_implementation, 

245 validate_placement, 

246 load_base_model, 

247 distributed_setup=None, 

248 qat_config=None, 

249 fp8_config=None, 

250 compile_config=None, 

251 freeze_config=None, 

252 activation_checkpoint=None, 

253 activation_checkpoint_selection=None, 

254 swap_inputs=False, 

255 allow_cache_entry_mutation=False, 

256 activation_swap="none", 

257 model_init_dtype=None, 

258 **kwargs, 

259 ) -> Any: 

260 """Construct and observe a meta model without runtime infrastructure.""" 

261 del ( 

262 fsdp2_manager, 

263 validate_placement, 

264 compile_config, 

265 activation_checkpoint, 

266 activation_checkpoint_selection, 

267 swap_inputs, 

268 allow_cache_entry_mutation, 

269 activation_swap, 

270 ) 

271 factory, context = auto_model._resolve_module_construction( # pylint: disable=protected-access 

272 kwargs, hf_config, distributed_setup, model_init_dtype or torch_dtype, 

273 ) 

274 init_kwargs = auto_model._model_init_kwargs( # pylint: disable=protected-access 

275 dict(kwargs), 

276 load_base_model=False, 

277 ) 

278 if "trust_remote_code" in kwargs: 

279 init_kwargs["trust_remote_code"] = kwargs["trust_remote_code"] 

280 _prepare_model_source_compat(context["model_spec"]) 

281 with _codegen_model_init_context(): 

282 if factory is not None: 

283 if model_args: 

284 raise ValueError("model_factory uses config and context, not additional positional model arguments") 

285 model = auto_model.construct_family_model(factory, hf_config, context=context) 

286 else: 

287 _, model = auto_model._init_model( # pylint: disable=protected-access 

288 cls, None, hf_config, attn_implementation, torch_dtype, is_hf_model, 

289 *model_args, **init_kwargs, 

290 ) 

291 auto_model.attach_build_context(model, context) 

292 _record_source( 

293 session, 

294 model, 

295 hf_config, 

296 pretrained_model_name_or_path, 

297 is_hf_model, 

298 ) 

299 model_builder._apply_pre_sharding_features( # pylint: disable=protected-access 

300 peft_config, 

301 qat_config, 

302 fp8_config, 

303 ) 

304 model = _observe_replacements( 

305 model, 

306 distributed_setup, 

307 getattr(distributed_setup, "low_precision_config", None), 

308 load_base_model, 

309 session, 

310 ) 

311 if freeze_config is not None: 

312 model_builder._apply_parameter_freezing( # pylint: disable=protected-access 

313 model, 

314 freeze_config, 

315 ) 

316 model_builder._apply_model_method_replacements( # pylint: disable=protected-access 

317 model, distributed_setup, getattr(distributed_setup, "low_precision_config", None), 

318 ) 

319 _observe_plan(model, mesh, sharding_planner, is_hf_model, session) 

320 snapshot = _publish_session(session, artifact_dir, result) 

321 result.snapshot = snapshot 

322 model._hp_codegen_snapshot = snapshot # pylint: disable=protected-access 

323 model.train() 

324 return model 

325 

326 return _build_model 

327 

328 

329@contextmanager 

330def use_codegen_model_build_overlay( 

331 session: CodegenBuildSession, 

332 artifact_dir: Optional[str] = None, 

333) -> Iterator[CodegenModelBuildResult]: 

334 """Temporarily replace AutoModel's final build step for offline codegen.""" 

335 model_class = auto_model._BaseHyperAutoModelClass # pylint: disable=protected-access 

336 original = model_class.__dict__["_build_model"] 

337 result = CodegenModelBuildResult() 

338 model_class._build_model = classmethod( # pylint: disable=protected-access 

339 _codegen_build_model(session, artifact_dir, result) 

340 ) 

341 try: 

342 yield result 

343 finally: 

344 model_class._build_model = original # pylint: disable=protected-access 

345 

346 

347__all__ = ["CodegenModelBuildResult", "use_codegen_model_build_overlay"]