Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / tensor_parallel / head_count.py: 94%

1# Copyright 2025-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"""TP-local head-count adjustment for non-TP-tolerant modeling code (D-17). 

16 

17Some HF modeling code reshapes/splits with an explicit (global) head count, 

18e.g. ``q.view(b, s, self.num_heads, self.head_dim)``, instead of the 

19TP-tolerant ``-1`` form. After TP colwise sharding each rank holds 

20``num_heads / tp`` heads locally, so any forward that runs on **local 

21tensors** needs the local count (AutoModel-style adaptation). 

22 

23Dual-mode rule -- a module's cached attributes are rewritten only when its 

24forward sees local tensors in the current mode: 

25 

26- production: parameters are permanently unwrapped to local tensors, so 

27 every head-sharded module is rewritten; 

28- validate: boundary modules run DTensor dispatch on the GLOBAL logical 

29 shape (a local count would break the reshape), so they are never 

30 rewritten; local-region modules (``local_compute_fn`` / 

31 ``region_dispatch=False``) unwrap to local inside the region in both 

32 modes, so they are rewritten in validate too. 

33 

34The config object is never touched (``head_dim`` / RoPE derivations keep 

35working); only module-instance cached attributes are rewritten, idempotently 

36-- originals are stashed in ``module._hp_full_head_counts`` and repeat calls 

37are no-ops. 

38""" 

39 

40import logging 

41from dataclasses import dataclass 

42from typing import Any, Optional, Sequence, Tuple 

43 

44from hyper_parallel.distributed.recipe_spec import resolve_placements 

45from hyper_parallel.core.dtensor.placement_types import Shard 

46 

47logger = logging.getLogger(__name__) 

48 

49 

50@dataclass(frozen=True) 

51class TpLocalAttrPlan: 

52 """Planner-generated TP-local module attribute adjustment plan.""" 

53 

54 auto_divide: Tuple[str, ...] = () 

55 user_divide: Tuple[str, ...] = () 

56 

57# Attribute list: survey of transformers/src/transformers/models (2026-07), 

58# counting forward-time reshape/split usages of cached head-count attributes: 

59# - q head count: num_heads (x393) / num_attention_heads (x122) / n_heads 

60# (x50) / num_attn_heads (prophetnet) / n_head (openai, xlnet) / 

61# heads, num_head (falcon, qwen2_5_omni, ...); 

62# - kv head count: num_key_value_heads (widespread) / num_kv_heads, kv_heads 

63# (falcon, gpt_bigcode); 

64# - excluded: head_dim / attention_head_size / head_size (the head dimension, 

65# never sharded) and num_key_value_groups (a ratio -- TP-invariant). 

66Q_HEAD_ATTRS = ( 

67 "num_heads", "num_attention_heads", "n_heads", "num_attn_heads", 

68 "n_head", "heads", "num_head", 

69) 

70KV_HEAD_ATTRS = ("num_key_value_heads", "num_kv_heads", "kv_heads") 

71 

72_ORIGINALS_ATTR = "_hp_full_head_counts" 

73_TP_LOCAL_ORIGINALS_ATTR = "_hp_full_tp_local_attrs" 

74 

75_FORBIDDEN_USER_ATTRS = frozenset({ 

76 "head_dim", "attention_head_size", "head_size", 

77 "num_key_value_groups", "training", "dtype", "device", 

78}) 

79 

80# q/k/v projections whose colwise Shard(0) splits the head dimension 

81# (q_b_proj covers the MLA up-projection, D-14). 

82_QKV_WEIGHT_SUFFIXES = ( 

83 "q_proj.weight", "q_b_proj.weight", "k_proj.weight", "v_proj.weight", 

84 "qkv_proj.weight", "qkv.weight", 

85) 

86 

87 

88def _is_plain_int(value: Any) -> bool: 

89 """Return whether a value is an integer but not a boolean.""" 

90 return isinstance(value, int) and not isinstance(value, bool) 

91 

92 

93def _is_head_sharded(spec, mesh_dim_names) -> bool: 

94 """True when the spec column-wise shards any q/k/v projection on the TP axis.""" 

95 if "tp" not in mesh_dim_names: 

96 return False 

97 tp_idx = tuple(mesh_dim_names).index("tp") 

98 for name, placement in spec.params.items(): 

99 if not name.endswith(_QKV_WEIGHT_SUFFIXES): 

100 continue 

101 if resolve_placements(placement, mesh_dim_names)[tp_idx] == Shard(0): 

102 return True 

103 return False 

104 

105 

106def _tp_degree(mesh, mesh_dim_names) -> int: 

107 if "tp" not in mesh_dim_names: 

108 return 1 

109 return mesh["tp"].size() 

110 

111 

112def update_module_head_counts(module: Any, tp_size: int, module_fqn: str = "") -> int: 

113 """Divide a module's cached head-count attributes by tp_size (idempotent). 

114 

115 Only plain-int attributes whose name is in Q_HEAD_ATTRS / KV_HEAD_ATTRS 

116 are rewritten. Originals are recorded in ``module._hp_full_head_counts``; 

117 repeat calls (e.g. applying a plan twice) are no-ops. A non-divisible 

118 value is left unchanged with a loud warning (planner-side 

119 validate_model_compatibility normally fails first on config-level 

120 divisibility). 

121 

122 Returns the number of attributes rewritten on this call. 

123 """ 

124 if tp_size <= 1: 

125 return 0 

126 originals = getattr(module, _ORIGINALS_ATTR, None) 

127 if originals is None: 

128 originals = {} 

129 setattr(module, _ORIGINALS_ATTR, originals) 

130 n = 0 

131 for attr in Q_HEAD_ATTRS + KV_HEAD_ATTRS: 

132 if attr in originals: 

133 continue 

134 value = getattr(module, attr, None) 

135 if not _is_plain_int(value): 

136 continue 

137 originals[attr] = value 

138 if value % tp_size != 0: 

139 logger.warning( 

140 "head-count adjustment: %s.%s = %d is not divisible by " 

141 "tp_size=%d -- left unchanged", 

142 module_fqn, attr, value, tp_size) 

143 continue 

144 setattr(module, attr, value // tp_size) 

145 logger.info("head-count adjustment: %s.%s %d -> %d (tp_size=%d)", 

146 module_fqn, attr, value, value // tp_size, tp_size) 

147 n += 1 

148 return n 

149 

150 

151def collect_auto_head_attrs( 

152 module: Any, spec: Any, mesh_dim_names: Sequence[str], 

153) -> tuple[str, ...]: 

154 """Collect canonical head-count attributes derived from TP placements.""" 

155 if not _is_head_sharded(spec, mesh_dim_names): 

156 return () 

157 return tuple( 

158 attr for attr in Q_HEAD_ATTRS + KV_HEAD_ATTRS 

159 if _is_plain_int(getattr(module, attr, None)) 

160 ) 

161 

162 

163def normalize_tp_divide_attrs( 

164 attrs: Optional[Sequence[str]], *, module: Any, module_fqn: str, 

165 tp_size: int, 

166) -> tuple[str, ...]: 

167 """Validate and normalize user-declared TP-local integer attributes.""" 

168 if attrs is None: 

169 return () 

170 if not isinstance(attrs, (list, tuple)): 

171 raise ValueError( 

172 f"{module_fqn}: tp_divide_attrs must be a list of attribute names, " 

173 f"got {type(attrs).__name__}") 

174 normalized = [] 

175 seen = set() 

176 for attr in attrs: 

177 if not isinstance(attr, str) or not attr or not attr.isidentifier(): 

178 raise ValueError( 

179 f"{module_fqn}: tp_divide_attrs entries must be non-empty " 

180 f"Python identifiers, got {attr!r}") 

181 if attr.startswith("_hp_") or attr in _FORBIDDEN_USER_ATTRS: 

182 raise ValueError( 

183 f"{module_fqn}: tp_divide_attrs cannot adjust protected " 

184 f"attribute {attr!r}") 

185 if attr in seen: 

186 raise ValueError( 

187 f"{module_fqn}: tp_divide_attrs contains duplicate {attr!r}") 

188 seen.add(attr) 

189 value = getattr(module, attr, None) 

190 if not _is_plain_int(value): 

191 actual = type(value).__name__ if hasattr(module, attr) else "missing" 

192 raise ValueError( 

193 f"{module_fqn}: tp_divide_attrs attribute {attr!r} must " 

194 f"exist and be a plain int, got {actual}") 

195 if value <= 0 or value % tp_size != 0: 

196 raise ValueError( 

197 f"{module_fqn}: tp_divide_attrs attribute {attr!r}={value} " 

198 f"must be positive and divisible by tp_size={tp_size}") 

199 normalized.append(attr) 

200 return tuple(normalized) 

201 

202 

203def build_tp_local_attr_plan( 

204 module: Any, spec: Any, module_fqn: str, tp_size: int, 

205 mesh_dim_names: Sequence[str], 

206) -> TpLocalAttrPlan: 

207 """Build the internal auto/user TP-local attribute plan for one module.""" 

208 auto_attrs = collect_auto_head_attrs(module, spec, mesh_dim_names) 

209 user_attrs = normalize_tp_divide_attrs( 

210 spec.tp_divide_attrs, module=module, module_fqn=module_fqn, 

211 tp_size=tp_size, 

212 ) 

213 duplicates = sorted(set(auto_attrs) & set(user_attrs)) 

214 if duplicates: 

215 raise ValueError( 

216 f"{module_fqn}: tp_divide_attrs redundantly declares D-17 " 

217 f"automatic head attributes {duplicates}; remove them from YAML") 

218 return TpLocalAttrPlan(auto_divide=auto_attrs, user_divide=user_attrs) 

219 

220 

221def _update_user_tp_attrs( 

222 module: Any, attrs: Sequence[str], tp_size: int, module_fqn: str, 

223) -> int: 

224 """Apply validated user TP-local attribute divisions idempotently.""" 

225 if tp_size <= 1 or not attrs: 

226 return 0 

227 originals = getattr(module, _TP_LOCAL_ORIGINALS_ATTR, None) 

228 if originals is None: 

229 originals = {} 

230 setattr(module, _TP_LOCAL_ORIGINALS_ATTR, originals) 

231 count = 0 

232 for attr in attrs: 

233 if attr in originals: 

234 expected = originals[attr] // tp_size 

235 current = getattr(module, attr) 

236 if current != expected: 

237 raise ValueError( 

238 f"{module_fqn}: tp_divide_attrs attribute {attr!r} was " 

239 f"already adjusted from {originals[attr]} to {current}, " 

240 f"which is incompatible with tp_size={tp_size} " 

241 f"(expected {expected})") 

242 continue 

243 value = getattr(module, attr) 

244 if not _is_plain_int(value): 

245 raise ValueError( 

246 f"{module_fqn}: tp_divide_attrs attribute {attr!r} must " 

247 f"exist and be a plain int, got {type(value).__name__}") 

248 if value <= 0 or value % tp_size != 0: 

249 raise ValueError( 

250 f"{module_fqn}: tp_divide_attrs attribute {attr!r}={value} " 

251 f"must be positive and divisible by tp_size={tp_size}") 

252 originals[attr] = value 

253 setattr(module, attr, value // tp_size) 

254 logger.info( 

255 "TP-local attribute adjustment: %s.%s %d -> %d (tp_size=%d, source=user)", 

256 module_fqn, attr, value, value // tp_size, tp_size, 

257 ) 

258 count += 1 

259 return count 

260 

261 

262def maybe_update_head_counts( 

263 module: Any, 

264 spec: Any, 

265 module_fqn: str, 

266 mesh: Any, 

267 mesh_dim_names: Sequence[str], 

268) -> None: 

269 """Adjust cached head counts when the spec head-shards the module. 

270 

271 Called by the applier exactly where the module's forward is guaranteed to 

272 see local tensors: production (Phase A, all specs) and validate (Phase C, 

273 local-region specs only). 

274 """ 

275 tp_size = _tp_degree(mesh, mesh_dim_names) 

276 attr_plan = getattr(spec, "_tp_local_attr_plan", None) 

277 if attr_plan is None: 

278 # Backward-compatible path for manually assembled plans that did not 

279 # pass through ShardingPlanner.finalize. 

280 if _is_head_sharded(spec, mesh_dim_names): 

281 update_module_head_counts(module, tp_size, module_fqn) 

282 _update_user_tp_attrs( 

283 module, getattr(spec, "tp_divide_attrs", ()) or (), 

284 tp_size, module_fqn, 

285 ) 

286 return 

287 if attr_plan.auto_divide: 

288 update_module_head_counts(module, tp_size, module_fqn) 

289 _update_user_tp_attrs( 

290 module, attr_plan.user_divide, tp_size, module_fqn, 

291 )