Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / plan / derive.py: 18%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# Licensed under the Apache License, Version 2.0 

3# ============================================================================ 

4"""Derive the sharding plan for a model without a live process group. 

5 

6The real ``ShardingPlanner`` runs at *generation* time so the 

7frozen plan (and the signature over it) reflects the actual model structure 

8— planner outputs land in ``CodegenMeta`` and, later, in the generated 

9modeling file. All helpers are deliberately independent of the training 

10stack: they take a ``CodegenSpec`` (or a config object exposing the same 

11attributes) and return plain data. 

12 

13``build_meta_model`` constructs the model on the meta device — no weights 

14are loaded, so a codegen run never needs the checkpoint or a GPU. 

15""" 

16from __future__ import annotations 

17 

18import logging 

19from typing import Any, Optional 

20 

21from hyper_parallel.codegen.plan.offline_mesh import build_offline_mesh 

22from hyper_parallel.codegen.spec.types import CodegenSpec 

23 

24logger = logging.getLogger(__name__) 

25 

26 

27def _custom_model_cls(spec: CodegenSpec, hf_config: Any) -> Optional[type]: 

28 """Resolve a custom-registered model class for the spec/config, if any.""" 

29 try: 

30 from hyper_parallel.models.registry import ( # pylint: disable=C0415 

31 _resolve_custom_model_cls, 

32 get_model_adapter, 

33 ) 

34 except ImportError: 

35 return None 

36 candidates = ( 

37 list(getattr(hf_config, "architectures", None) or []) 

38 if hf_config is not None else [] 

39 ) 

40 spec_arch = getattr(getattr(spec, "source", None), "architecture", None) 

41 if spec_arch: 

42 candidates.append(spec_arch) 

43 for arch in candidates: 

44 if not arch: 

45 continue 

46 if _resolve_custom_model_cls(arch) is None: 

47 # Lazy family discovery: either spelling may be the key that 

48 # registers the family (mirrors get_is_hf_model's identity loop). 

49 identities = [arch] 

50 if hf_config is not None: 

51 identities.append(getattr(hf_config, "model_type", "") or "") 

52 for identity in identities: 

53 if identity: 

54 get_model_adapter(identity) 

55 cls = _resolve_custom_model_cls(arch) 

56 if cls is not None: 

57 return cls 

58 return None 

59 

60 

61def build_meta_model( 

62 spec: CodegenSpec, 

63 *, 

64 config_overrides: Optional[dict[str, Any]] = None, 

65 hf_config: Any = None, 

66) -> Any: 

67 """Construct an empty meta-device model for the spec's source. 

68 

69 ``AutoModelForCausalLM.from_config(config)`` builds the full parameter 

70 tree with no weights and no network; wrapping it in 

71 ``components.utils.model_utils.init_empty_weights`` keeps the parameters 

72 on the meta device so a CPU-only host never allocates the model's memory. 

73 

74 Returns ``None`` only when the spec carries no model identity (the caller 

75 then skips plan derivation entirely). A model that *is* configured but 

76 cannot be built is a hard failure — swallowing it would freeze an empty 

77 plan and let an incomplete artifact pass preflight. 

78 """ 

79 model_id = spec.source.model_name_or_path 

80 if not model_id: 

81 return None 

82 try: 

83 from hyper_parallel import init_empty_weights 

84 

85 custom_cls = _custom_model_cls(spec, hf_config) 

86 if custom_cls is not None: 

87 # A custom-registered architecture is unknown to transformers 

88 # AutoModel; build it the way the trainer's CUSTOM/GEN backends 

89 # do — from the resolved config, which also carries any crop. 

90 if hf_config is None: 

91 raise ValueError( 

92 "codegen: custom-registered model " 

93 f"{spec.source.architecture or model_id!r} needs the " 

94 "resolved config (hf_config) to build the meta model; " 

95 "AutoConfig cannot load it from the model id" 

96 ) 

97 with init_empty_weights(): 

98 return custom_cls(hf_config) 

99 

100 from transformers import AutoConfig, AutoModelForCausalLM 

101 

102 config = AutoConfig.from_pretrained(model_id, **dict(config_overrides or {})) 

103 with init_empty_weights(): 

104 return AutoModelForCausalLM.from_config(config) 

105 except Exception as exc: # noqa: BLE001 - fail fast: an unbuildable source must not yield a plan-less artifact 

106 logger.error( 

107 "codegen: cannot build meta model for %r (%s: %s)", 

108 model_id, type(exc).__name__, exc, 

109 ) 

110 raise 

111 

112 

113def build_plan_overrides(config: Any, spec: CodegenSpec) -> dict[str, Any]: 

114 """Desugar ``config.plan_overrides`` into the planner's override dict. 

115 

116 Wraps ``trainer.config.entries_to_plan_overrides`` (imported lazily — the 

117 codegen package must stay importable without the trainer layer). 

118 ``when``-gated entries are skipped here against the projected dims, the 

119 same way the trainer's own desugar pass applies them at run time. 

120 """ 

121 entries = getattr(config, "plan_overrides", None) or [] 

122 if not entries: 

123 return {} 

124 from hyper_parallel.trainer.config import entries_to_plan_overrides 

125 

126 return entries_to_plan_overrides( 

127 entries, 

128 cp_size=spec.parallel_dims.cp_size, 

129 ep_size=spec.parallel_dims.ep_size, 

130 ) 

131 

132 

133def derive_sharding_plan( 

134 model: Any, spec: CodegenSpec, *, plan_overrides: Optional[dict[str, Any]] = None, 

135) -> Any: 

136 """Run ``ShardingPlanner.plan`` over the offline mesh for ``spec``. 

137 

138 Returns the planner's ``ShardingPlan``. A planner failure raises — the 

139 manager's ``_fill_plan_fields`` turns a None plan into an empty frozen 

140 plan, which is exactly the silent-degradation path that must not survive 

141 (a plan-less artifact passing preflight). ``model`` is the meta-device 

142 model from :func:`build_meta_model`; the mesh is the static offline mesh, 

143 so no process group is needed. 

144 

145 ``plan_overrides`` is the desugared override dict — the planner's native 

146 interface — normally produced by :func:`build_plan_overrides` from the 

147 live config. ``allow_uncovered_params`` stays at its default (False): 

148 the F4b fail-fast guard must fire here, at generation time, not silently 

149 inside the trainer. 

150 """ 

151 parallel_dims = spec.parallel_dims 

152 mesh = build_offline_mesh(parallel_dims) 

153 from hyper_parallel.distributed._builder.planner import ShardingPlanner 

154 

155 planner = ShardingPlanner(plan_overrides=dict(plan_overrides or {})) 

156 return planner.plan( 

157 model, 

158 mesh, 

159 tp_size=parallel_dims.tp_size, 

160 cp_size=parallel_dims.cp_size, 

161 ep_size=parallel_dims.ep_size, 

162 sequence_parallel=parallel_dims.sequence_parallel, 

163 loss_parallel=parallel_dims.loss_parallel, 

164 ) 

165 

166 

167__all__ = [ 

168 "build_meta_model", 

169 "build_plan_overrides", 

170 "derive_sharding_plan", 

171]