Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed_data / cost_model.py: 71%

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"""Replaceable workload estimates, independent of packing hard limits.""" 

16 

17from __future__ import annotations 

18 

19import math 

20from collections.abc import Mapping 

21from dataclasses import dataclass 

22from typing import Any, Protocol 

23 

24from hyper_parallel.distributed_data.schema import SampleMetadata, WorkloadCost 

25 

26 

27class CostModel(Protocol): 

28 """Deterministic, stateless CPU estimate evaluated on the centralized planner. 

29 

30 Implementations consume lightweight ``metadata.features``, not image tensors. 

31 They must not mutate metadata, perform distributed collectives, or change 

32 ``pack_tokens`` / ``packing_costs``. Return components in comparable units. 

33 Models may also expose ``padding_cost(padding_tokens) -> WorkloadCost`` for 

34 constructor padding. The planner and built-in LPT policies include it in 

35 each pack's layout-dependent cost; custom policies still receive source costs. 

36 """ 

37 

38 def __call__(self, metadata: SampleMetadata) -> WorkloadCost: 

39 """Estimate one selected sample's additive workload.""" 

40 

41 

42class CausalAttentionCostModel: 

43 """Normalized causal-attention square cost for one independent text document. 

44 

45 Costs use token-squared units, omitting common model coefficients and all 

46 linear terms. The constructor's independent PAD sequence contributes its 

47 own squared length through ``padding_cost`` after each candidate placement. 

48 """ 

49 

50 def __call__(self, metadata: SampleMetadata) -> WorkloadCost: 

51 """Return L squared without counting ignored targets as padding.""" 

52 if metadata.features.get("D", 0): 

53 raise ValueError("CausalAttentionCostModel requires causal text documents without diffusion tokens") 

54 return WorkloadCost(llm=metadata.pack_tokens ** 2) 

55 

56 @staticmethod 

57 def padding_cost(padding_tokens: int) -> WorkloadCost: 

58 """Account for the independent PAD sequence that Attention still computes. 

59 

60 Args: 

61 padding_tokens: Number of tokens appended to one packed row. 

62 

63 Returns: 

64 Normalized causal Attention square cost of that PAD sequence. 

65 """ 

66 return WorkloadCost(llm=padding_tokens ** 2) 

67 

68 

69def _configuration_dict(config: Any) -> dict[str, Any]: 

70 """Accept an ordinary dictionary or a configuration with ``to_dict``.""" 

71 if isinstance(config, Mapping): 

72 return dict(config) 

73 if callable(getattr(config, "to_dict", None)): 

74 value = config.to_dict() 

75 if isinstance(value, Mapping): 

76 return dict(value) 

77 raise ValueError("model_config must be a mapping or a configuration exposing to_dict().") 

78 

79 

80def _integer(value: Any, name: str, minimum: int = 1) -> int: 

81 if not isinstance(value, int) or isinstance(value, bool) or value < minimum: 

82 raise ValueError(f"{name} must be an integer >= {minimum}, but got {value!r}.") 

83 return value 

84 

85 

86@dataclass(frozen=True) 

87class BackboneFlopsConfig: 

88 """Latent-attention decoder dimensions, independent of model libraries. 

89 

90 Only the shared backbone decoder layers are modeled. Diffusion input/output 

91 stacks, the language-model head, ViT, VAE, and their projectors are excluded. 

92 ``mlp_layer_types`` follows the actual layer order: ``dense`` or ``sparse``. 

93 """ 

94 

95 hidden_size: int 

96 num_hidden_layers: int 

97 num_attention_heads: int 

98 intermediate_size: int 

99 mlp_layer_types: tuple[str, ...] 

100 kv_lora_rank: int 

101 q_lora_rank: int | None 

102 qk_nope_head_dim: int 

103 qk_rope_head_dim: int 

104 v_head_dim: int 

105 moe_intermediate_size: int = 0 

106 n_routed_experts: int = 0 

107 num_experts_per_tok: int = 0 

108 n_shared_experts: int = 0 

109 attn_block_mode: str = "tail" 

110 diff_time_embed_mode: str = "add" 

111 

112 def __post_init__(self) -> None: 

113 """Reject dimensions or layer patterns not covered by this estimate.""" 

114 for name in ( 

115 "hidden_size", "num_hidden_layers", "num_attention_heads", "intermediate_size", 

116 "kv_lora_rank", "v_head_dim", 

117 ): 

118 _integer(getattr(self, name), name) 

119 for name in ("qk_nope_head_dim", "qk_rope_head_dim", "n_shared_experts"): 

120 _integer(getattr(self, name), name, 0) 

121 if self.q_lora_rank is not None: 

122 _integer(self.q_lora_rank, "q_lora_rank") 

123 if self.qk_nope_head_dim + self.qk_rope_head_dim == 0: 

124 raise ValueError("The combined query/key head dimension must be positive.") 

125 if ( 

126 not isinstance(self.mlp_layer_types, tuple) 

127 or len(self.mlp_layer_types) != self.num_hidden_layers 

128 or any(kind not in ("dense", "sparse") for kind in self.mlp_layer_types) 

129 ): 

130 raise ValueError("mlp_layer_types must contain one 'dense' or 'sparse' entry per backbone layer.") 

131 for name in ("moe_intermediate_size", "n_routed_experts", "num_experts_per_tok"): 

132 _integer(getattr(self, name), name, int("sparse" in self.mlp_layer_types)) 

133 if self.num_experts_per_tok > self.n_routed_experts: 

134 raise ValueError("num_experts_per_tok cannot exceed n_routed_experts.") 

135 if self.attn_block_mode not in ("tail", "blockmask"): 

136 raise ValueError("attn_block_mode must be 'tail' or 'blockmask'.") 

137 if self.diff_time_embed_mode not in ("add", "prepend", "adaln"): 

138 raise ValueError("diff_time_embed_mode must be 'add', 'prepend', or 'adaln'.") 

139 

140 @classmethod 

141 def from_model_config(cls, model_config: Any) -> "BackboneFlopsConfig": 

142 """Read the final model configuration after architecture overrides. 

143 

144 Args: 

145 model_config: Configuration with explicit decoder dimensions and layer types. 

146 """ 

147 config = _configuration_dict(model_config) 

148 text = _configuration_dict(config.get("text_config", config)) 

149 required = ( 

150 "hidden_size", "num_hidden_layers", "num_attention_heads", "intermediate_size", 

151 "kv_lora_rank", "q_lora_rank", "qk_nope_head_dim", "qk_rope_head_dim", "v_head_dim", 

152 ) 

153 missing = [name for name in required if name not in text] 

154 if missing: 

155 raise ValueError(f"Backbone cost model is missing configuration fields: {missing}.") 

156 pattern = text.get("mlp_layer_types") 

157 if not isinstance(pattern, (list, tuple)): 

158 raise ValueError("model_config must provide mlp_layer_types in actual decoder layer order.") 

159 qk_dim = ( 

160 _integer(text["qk_nope_head_dim"], "qk_nope_head_dim", 0) 

161 + _integer(text["qk_rope_head_dim"], "qk_rope_head_dim", 0) 

162 ) 

163 if text.get("qk_head_dim", qk_dim) != qk_dim: 

164 raise ValueError("qk_head_dim must equal qk_nope_head_dim + qk_rope_head_dim.") 

165 return cls( 

166 **{name: text[name] for name in required}, 

167 mlp_layer_types=tuple(pattern), 

168 **{name: text.get(name, 0) for name in ( 

169 "moe_intermediate_size", "n_routed_experts", "num_experts_per_tok", "n_shared_experts", 

170 )}, 

171 attn_block_mode=config.get("attn_block_mode", "tail"), 

172 diff_time_embed_mode=config.get("diff_time_embed_mode", "add"), 

173 ) 

174 

175 

176class DefaultCostModel: 

177 """Shape-based FLOPs for a shared latent-attention and expert backbone. 

178 

179 Args: 

180 model_config: Final model configuration or :class:`BackboneFlopsConfig`. 

181 Required to derive the backbone's arithmetic workload. 

182 training_multiplier: Constant forward-to-training FLOPs proxy: 1 for 

183 forward, 3 for forward/backward, or 4 for full-layer recomputation. 

184 It does not change balancing decisions when shared by all samples. 

185 

186 Note: 

187 MAC=2 FLOPs. The attention area per independent raw item is 

188 ``P**2 / 2 + D * (P + D)``; blockwise attention adds half the square of each 

189 conditional-image placeholder run. ``P`` includes control tokens and 

190 conditional-image tokens, not just natural-language text. Costs add 

191 across packed items; there is no attention between different items. 

192 

193 Linear terms include MLA Q/KV/O projections, dense SwiGLU, routed top-k 

194 and shared experts, and router projection. Norms, activations, softmax, 

195 attention tile padding, communication and expert-kernel efficiency are 

196 excluded. This is theoretical arithmetic, not a latency predictor. 

197 """ 

198 

199 def __init__(self, model_config: Any, *, training_multiplier: float = 3.0) -> None: 

200 """Compile workload coefficients from the effective decoder configuration. 

201 

202 Args: 

203 model_config: Explicit decoder dimensions and attention layout. 

204 training_multiplier: Forward-to-training arithmetic scale. 

205 """ 

206 if model_config is None: 

207 raise ValueError("DefaultCostModel requires model_config; provide it or supply a custom cost_model.") 

208 if ( 

209 not isinstance(training_multiplier, (int, float)) 

210 or isinstance(training_multiplier, bool) 

211 or not math.isfinite(training_multiplier) 

212 or training_multiplier <= 0 

213 ): 

214 raise ValueError("training_multiplier must be finite and positive.") 

215 self.training_multiplier = float(training_multiplier) 

216 self.config = ( 

217 model_config if isinstance(model_config, BackboneFlopsConfig) 

218 else BackboneFlopsConfig.from_model_config(model_config) 

219 ) 

220 self.linear_flops_per_token = 0 

221 self.attention_flops_per_pair = 0 

222 self._compile_coefficients() 

223 

224 def _compile_coefficients(self) -> None: 

225 """Precompute linear-token and attention-pair FLOPs from backbone dimensions.""" 

226 config = self.config 

227 hidden, heads = config.hidden_size, config.num_attention_heads 

228 qk_dim = config.qk_nope_head_dim + config.qk_rope_head_dim 

229 query_mac = hidden * heads * qk_dim if config.q_lora_rank is None else ( 

230 hidden * config.q_lora_rank + config.q_lora_rank * heads * qk_dim 

231 ) 

232 projection_mac = ( 

233 query_mac 

234 + hidden * (config.kv_lora_rank + config.qk_rope_head_dim) 

235 + config.kv_lora_rank * heads * (config.qk_nope_head_dim + config.v_head_dim) 

236 + heads * config.v_head_dim * hidden 

237 ) 

238 dense_layers = config.mlp_layer_types.count("dense") 

239 sparse_layers = config.num_hidden_layers - dense_layers 

240 dense_mac = 3 * hidden * config.intermediate_size 

241 sparse_mac = ( 

242 3 * hidden * config.moe_intermediate_size * (config.num_experts_per_tok + config.n_shared_experts) 

243 + hidden * config.n_routed_experts 

244 ) 

245 self.linear_flops_per_token = 2 * ( 

246 config.num_hidden_layers * projection_mac + dense_layers * dense_mac + sparse_layers * sparse_mac 

247 ) 

248 self.attention_flops_per_pair = 2 * config.num_hidden_layers * heads * (qk_dim + config.v_head_dim) 

249 

250 def forward_flops(self, metadata: SampleMetadata) -> float: 

251 """Estimate decoder forward FLOPs from CPU-only per-raw-item metadata. 

252 

253 Args: 

254 metadata: One independent item's token counts and attention segments. 

255 """ 

256 prefix = _integer(metadata.features.get("P"), "metadata.features['P']", 0) 

257 diffusion = _integer(metadata.features.get("D"), "metadata.features['D']", 0) 

258 if prefix + diffusion != metadata.pack_tokens: 

259 raise ValueError("metadata features P + D must equal pack_tokens before optional timestep-token insertion.") 

260 if diffusion and self.config.diff_time_embed_mode == "prepend": 

261 diffusion += 1 

262 pairs = 0.5 * prefix * prefix + diffusion * (prefix + diffusion) 

263 if self.config.attn_block_mode == "blockmask": 

264 runs = metadata.features.get("cond_image_token_lengths") 

265 if not isinstance(runs, (list, tuple)): 

266 raise ValueError("Blockwise attention requires cond_image_token_lengths (empty for no images).") 

267 lengths = [_integer(length, "conditional-image run length") for length in runs] 

268 if sum(lengths) > prefix: 

269 raise ValueError("Conditional-image run lengths cannot exceed the prefix token count P.") 

270 pairs += 0.5 * sum(length * length for length in lengths) 

271 return float(self.linear_flops_per_token * (prefix + diffusion) + self.attention_flops_per_pair * pairs) 

272 

273 def __call__(self, metadata: SampleMetadata) -> WorkloadCost: 

274 """Return backbone-only workload without changing physical packing caps.""" 

275 return WorkloadCost(llm=self.training_multiplier * self.forward_flops(metadata)) 

276 

277 

278def resolve_cost_model(cost_model: CostModel | None, model_config: Any = None) -> CostModel: 

279 """Use an explicit callback or construct the default from model dimensions. 

280 

281 Args: 

282 cost_model: Optional user-supplied workload callback. 

283 model_config: Required backbone configuration when the callback is absent. 

284 """ 

285 if cost_model is None: 

286 return DefaultCostModel(model_config) 

287 if not callable(cost_model): 

288 raise ValueError("cost_model must be callable: SampleMetadata -> WorkloadCost.") 

289 return cost_model 

290 

291 

292__all__ = ["BackboneFlopsConfig", "CostModel", "DefaultCostModel"]