Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/distributed/_builder/family

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

7# http://www.apache.org/licenses/LICENSE-2.0

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"""Discover complete module-owned layouts before generic parameter templates."""

17from __future__ import annotations

19import copy

20import fnmatch

21import inspect

22from dataclasses import dataclass, field

23from typing import Any

25from torch import nn

27from hyper_parallel.components.modules.module_family import (

28 collect_family_modules, _family_parameter_shapes, get_family_build_info,

29 _record_family_layouts, _validate_applied_layout,

30)

31from hyper_parallel.core.dtensor.placement_types import Partial

32from hyper_parallel.distributed.module_binding import (

33 NativeParallelBindingRequest, ParallelModuleSpec, validate_materialized_modules, validate_parallel_module,

34)

35from hyper_parallel.distributed.recipe_spec import ModuleShardingSpec, _normalize_out_fields

38@dataclass

39class OwnedLayouts:

40 """Construction-local ownership, indexed by final Module/Parameter identity."""

42 specs: dict[str, ModuleShardingSpec] = field(default_factory=dict)

43 parallel: dict[str, ParallelModuleSpec] = field(default_factory=dict)

44 owned_names: set[str] = field(default_factory=set)

45 expected_ids: set[int] = field(default_factory=set)

46 required_by_fqn: dict[str, set[int]] = field(default_factory=dict)

47 aliases: list[tuple[str, str]] = field(default_factory=list)

50def _full(fqn, name):

51 return f"{fqn}.{name}" if fqn else name

54def _is_child(fqn, parent):

55 return fqn.startswith(parent + ".") if parent else bool(fqn)

58def _module_parallel(module, member):

59 provider = getattr(module, "parallel_module_spec", None)

60 declaration = provider() if provider is not None else None

61 if declaration is not None and not isinstance(declaration, ParallelModuleSpec):

62 raise TypeError("parallel_module_spec() must return ParallelModuleSpec")

63 if member is not None and member.selectors["parallel"] and declaration is None:

64 raise ValueError("a parallel family member must declare its explicit execution protocol")

65 return declaration

68def _add_provider(entries, model, fqn, sharding, parallel):

69 module = model.get_submodule(fqn)

70 identity = id(module)

71 if identity in entries:

72 raise ValueError(f"multiple layout/communication owners for module {fqn!r}")

73 if not isinstance(sharding, ModuleShardingSpec):

74 raise TypeError("layout providers must return ModuleShardingSpec")

75 sharding = copy.deepcopy(sharding)

76 _normalize_out_fields(sharding)

77 entries[identity] = (fqn, module, sharding, parallel)

80def _native_providers(entries, model, spec, context, boundaries):

81 """Locate model-owned sharding providers for native parallel components."""

82 if spec is None or spec.native_parallel is None:

83 return

84 requests = spec.native_parallel(model=model, context=context)

85 all_modules = tuple(model.named_modules(remove_duplicate=False))

86 family_subtrees = {id(child) for boundary in boundaries for child in boundary.module.modules()}

87 for request in requests:

88 if not isinstance(request, NativeParallelBindingRequest):

89 raise TypeError("native_parallel must return NativeParallelBindingRequest objects")

90 matched = set()

91 for pattern in request.match:

92 hits = [(fqn, module) for fqn, module in all_modules if fnmatch.fnmatchcase(fqn, pattern)]

93 if not hits:

94 raise ValueError(f"native binding pattern {pattern!r} matched no module")

95 for fqn, module in hits:

96 if id(module) in matched:

97 continue

98 if _overlaps_family(module, family_subtrees):

99 raise ValueError(f"native binding at {fqn!r} overlaps a family-owned subtree")

100 matched.add(id(module))

101 layout = request.sharding(module=module, module_fqn=fqn, member=None, context=context)

102 _add_provider(entries, model, fqn, layout, request.parallel)

103

104

105def _overlaps_family(module, family_subtrees):

106 return any(id(child) in family_subtrees for child in module.modules())

107

108

109def _validate_ports(module, spec):

110 """Validate named input and output layout declarations against the module interface."""

111 injection_fields = ("local_compute_fn", "inner_wrapper", "inner_target", "region_dispatch")

112 if any(getattr(spec, name) is not None for name in injection_fields):

113 raise ValueError("module-owned layouts cannot install legacy compute injection")

114 if spec.params is None or spec.in_src is None or spec.in_dst is None:

115 raise ValueError("module-owned providers must declare params, in_src and in_dst explicitly")

116 if set(spec.in_src) != set(spec.in_dst):

117 raise ValueError("module-owned in_src and in_dst must declare the same input names")

118 signature = inspect.signature(module.forward)

119 accepts_kwargs = any(item.kind == inspect.Parameter.VAR_KEYWORD for item in signature.parameters.values())

120 unknown = set(spec.in_src) - set(signature.parameters)

121 if unknown and not accepts_kwargs:

122 raise ValueError(f"layout declares unknown forward inputs {sorted(unknown)}")

123 if spec.out_names is not None and len(set(spec.out_names)) != len(spec.out_names):

124 raise ValueError("output names must be unique")

125

126

127def _validate_placements(spec, context):

128 """Reject unknown mesh axes and incompatible parameter placement declarations."""

129 mappings = [*spec.params.values(), *spec.in_src.values(), *spec.in_dst.values()]

130 mappings.extend((spec.out_src or {}).values())

131 mappings.extend((spec.out_dst or {}).values())

132 for named in mappings:

133 for axis, placement in named.items():

134 if axis not in context.axis_sizes:

135 raise ValueError(f"module layouts cannot declare storage/unknown axis {axis!r}")

136 if not hasattr(placement, "is_shard"):

137 raise TypeError("module providers must return concrete Placement objects")

138 if any(isinstance(placement, Partial) for named in spec.params.values() for placement in named.values()):

139 raise ValueError("parameter source layouts cannot be Partial")

140

141

142def _claim_layout_parameters(fqn, parameters, spec, children, owner_by_id, canonical_by_id):

143 """Assign each Parameter identity once and verify every alias has the same layout."""

144 declared = set()

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

146 parameter = parameters.get(name)

147 if parameter is None:

148 raise ValueError(f"layout parameter {_full(fqn, name)!r} does not exist")

149 full_name = _full(fqn, name)

150 if any(full_name.startswith(child + ".") for child in children):

151 raise ValueError(f"parent layout duplicates nested module ownership at {full_name!r}")

152 identity = id(parameter)

153 declared.add(identity)

154 if identity in owner_by_id:

155 if owner_by_id[identity] != placement:

156 raise ValueError(f"shared parameter {full_name!r} has conflicting layouts")

157 del spec.params[name]

158 else:

159 owner_by_id[identity] = placement

160 canonical_by_id[identity] = full_name

161 return declared

162

163

164def _assign_parameters(result, entries, model):

165 """Assign one owned layout per logical parameter while retaining shared aliases."""

166 all_parameters = dict(model.named_parameters(remove_duplicate=False))

167 owner_by_id = {}

168 canonical_by_id = {}

169 for fqn, module, spec, _ in entries.values():

170 parameters = dict(module.named_parameters(remove_duplicate=False))

171 children = [name for name, *_ in entries.values() if _is_child(name, fqn)]

172 expected = {id(parameter) for name, parameter in parameters.items()

173 if not any(_full(fqn, name).startswith(child + ".") for child in children)}

174 declared = _claim_layout_parameters(fqn, parameters, spec, children, owner_by_id, canonical_by_id)

175 if expected - declared:

176 missing = [name for name, parameter in parameters.items() if id(parameter) in expected - declared]

177 raise ValueError(f"module-owned layout {fqn!r} omits parameters {missing}")

178 result.expected_ids.update(expected)

179 result.specs[fqn] = spec

180 result.required_by_fqn[fqn] = {id(parameters[name]) for name in spec.params}

181 for name, parameter in all_parameters.items():

182 identity = id(parameter)

183 if identity not in owner_by_id:

184 continue

185 result.owned_names.add(name)

186 canonical = canonical_by_id[identity]

187 if canonical != name:

188 result.aliases.append((canonical, name))

189

190

191def collect_owned_layouts(model: nn.Module, context: Any, spec: Any = None) -> OwnedLayouts:

192 """Discover model/native/family providers independently of generic TP templates.

193

194 Args:

195 model: Final logical model, before any parameter sharding.

196 context: FamilyShardingContext with logical axes and global parameter shapes.

197 spec: Effective per-build model specification, if any.

198 """

199 entries = {}

200 if spec is not None and spec.sharding is not None:

201 for fqn, layout in spec.sharding(model=model, context=context).items():

202 _add_provider(entries, model, fqn, layout, _module_parallel(model.get_submodule(fqn), None))

203 boundaries = collect_family_modules(model)

204 for boundary in boundaries:

205 member = boundary.info.resolved

206 if member.sharding is None:

207 continue

208 fqn = boundary.fqns[0]

209 cached = boundary.info.layout

210 layout = (cached.sharding if cached is not None else

211 member.sharding(module=boundary.module, module_fqn=fqn, member=member, context=context))

212 entries.pop(id(boundary.module), None)

213 parallel = cached.parallel if cached is not None else _module_parallel(boundary.module, member)

214 _add_provider(entries, model, fqn, layout, parallel)

215 _native_providers(entries, model, spec, context, boundaries)

216 _validate_parallel_nesting(entries)

217 result = OwnedLayouts()

218 for fqn, module, layout, parallel in entries.values():

219 _validate_ports(module, layout)

220 _validate_placements(layout, context)

221 info = get_family_build_info(module)

222 cached = info.layout if info is not None else None

223 if parallel is not None and cached is None:

224 validate_parallel_module(module, parallel, layout, context)

225 if parallel is not None:

226 result.parallel[fqn] = parallel

227 _assign_parameters(result, entries, model)

228 return result

229

230

231def _validate_parallel_nesting(entries):

232 active = [(fqn, parallel.binding_kind) for fqn, _, _, parallel in entries.values()

233 if parallel is not None and parallel.binding_kind != "none"]

234 for fqn, kind in active:

235 for parent, parent_kind in active:

236 if _is_child(fqn, parent) and (kind == parent_kind or "mesh" in (kind, parent_kind)):

237 raise ValueError(f"overlapping communication owners at {parent!r} and {fqn!r}")

238

239

240def validate_owned_layouts(model: nn.Module, plan: Any, owned: OwnedLayouts, context: Any) -> None:

241 """Check ownership and member capabilities after the final explicit overrides.

242

243 Args:

244 model: Final model tree, still carrying global logical shapes.

245 plan: Plan after user layout overrides.

246 owned: Ownership captured before templates.

247 context: Shared immutable logical layout context.

248 """

249 if not owned.specs:

250 return

251 covered = set()

252 declarations = {}

253 for fqn, spec in plan.modules.items():

254 module = model.get_submodule(fqn)

255 parameters = dict(module.named_parameters(remove_duplicate=False))

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

257 parameter = parameters.get(name)

258 if parameter is None:

259 continue

260 identity = id(parameter)

261 if identity in declarations and declarations[identity] != placement:

262 raise ValueError(f"explicit overrides give shared parameter {_full(fqn, name)!r} conflicting layouts")

263 declarations[identity] = placement

264 covered.add(identity)

265 if fqn in owned.specs:

266 local_coverage = {id(parameters[name]) for name in spec.params if name in parameters}

267 if owned.required_by_fqn[fqn] - local_coverage:

268 raise ValueError(f"final override removed required parameter ownership from {fqn!r}")

269 _validate_ports(module, spec)

270 _validate_placements(spec, context)

271 parallel = owned.parallel.get(fqn)

272 info = get_family_build_info(module)

273 cached = info.layout if info is not None else None

274 if cached is not None:

275 _validate_applied_layout(cached, spec, fqn)

276 elif parallel is not None:

277 validate_parallel_module(module, parallel, spec, context)

278 if owned.expected_ids - covered:

279 raise ValueError("final layout overrides removed required module-owned parameter coverage")

280

281

282def _prepare_family_layout(module: nn.Module, mesh: Any) -> None:

283 """Apply a self-described component through the existing planner and executor."""

284 # Planner and apply import layout discovery; defer these cyclic runtime dependencies.

285 from hyper_parallel.distributed._builder.planner import ShardingPlanner # pylint: disable=C0415,R0401

286 from hyper_parallel.distributed.apply import apply_sharding_plan # pylint: disable=C0415,R0401

287

288 info = get_family_build_info(module)

289 if info.resolved.sharding is None:

290 if info.resolved.selectors["parallel"]:

291 raise ValueError("a mesh-constructed parallel family must declare sharding")

292 return

293 # A component's config identifies its containing model, whose providers do not apply here.

294 module._hp_model_adapter_spec = None # pylint: disable=protected-access

295 planner = ShardingPlanner()

296 plan = planner.plan(module, mesh.device_mesh, tp_size=mesh.tp_size, cp_size=mesh.cp_size,

297 ep_size=mesh.ep_size, sequence_parallel=mesh.sequence_parallel,

298 loss_parallel=mesh.loss_parallel)

299 for fqn, spec in plan.modules.items():

300 if fqn not in plan.parallel_modules:

301 validate_parallel_module(module.get_submodule(fqn), ParallelModuleSpec("family.local"), spec, None)

302 _record_family_layouts(module, plan, mesh, _family_parameter_shapes(module), None, applied=False)

303 if callable(getattr(mesh.device_mesh, "get_group", None)):

304 apply_sharding_plan(module, plan, mesh)

305 if not any(parameter.is_meta for parameter in module.parameters()):

306 validate_materialized_modules(module)

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / _builder / family_layout.py: 89%

220 statements