Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / module_binding.py: 78%

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"""Explicit stage protocols and resource-only binding for family and native modules.""" 

16 

17from __future__ import annotations 

18 

19import inspect 

20from copy import deepcopy 

21from collections.abc import Callable, Mapping 

22from dataclasses import dataclass, field 

23from types import MappingProxyType 

24from typing import Any 

25 

26from torch import nn 

27 

28from hyper_parallel.components.modules.module_family import _same_mesh, get_family_build_info 

29from hyper_parallel.core.dtensor.placement_types import Shard 

30from hyper_parallel.distributed.context_parallel.cp_context import CPContext 

31from hyper_parallel.distributed.expert_parallel.ep_context import EPContext 

32 

33 

34@dataclass(frozen=True) 

35class ParallelModuleSpec: 

36 """A class-owned execution protocol; validation proves its supported layouts.""" 

37 

38 protocol: str 

39 binding_kind: str = "none" 

40 stages: Mapping[str, tuple[str, ...]] = field(default_factory=dict) 

41 validate: Callable | None = None 

42 options: Mapping[str, Any] = field(default_factory=dict) 

43 version: str = "1" 

44 validate_materialized: Callable | None = None 

45 

46 def __post_init__(self) -> None: 

47 if not self.protocol or not self.version: 

48 raise ValueError("parallel protocol and version must be non-empty") 

49 if self.binding_kind not in ("none", "cp", "ep", "mesh"): 

50 raise ValueError("binding_kind must be none, cp, ep or mesh") 

51 if self.binding_kind != "none" and (not self.stages or self.validate is None): 

52 raise ValueError("parallel members must declare stage signatures and a layout validator") 

53 object.__setattr__(self, "stages", MappingProxyType({name: tuple(args) for name, args in self.stages.items()})) 

54 object.__setattr__(self, "options", MappingProxyType(deepcopy(dict(self.options)))) 

55 

56 

57@dataclass(frozen=True) 

58class NativeParallelBindingRequest: 

59 """Bind an existing component whose class already delegates explicit stages.""" 

60 

61 match: tuple[str, ...] 

62 parallel: ParallelModuleSpec 

63 sharding: Callable 

64 

65 def __post_init__(self) -> None: 

66 if isinstance(self.match, str) or not self.match or not all(isinstance(name, str) for name in self.match): 

67 raise ValueError("native matches must be a non-empty tuple of module FQN patterns") 

68 if not callable(self.sharding) or not isinstance(self.parallel, ParallelModuleSpec): 

69 raise ValueError("native binding requires a complete layout provider and ParallelModuleSpec") 

70 object.__setattr__(self, "match", tuple(self.match)) 

71 

72 

73@dataclass(frozen=True) 

74class MeshBinding: 

75 """Borrowed topology; all communication groups are owned by MeshContext.""" 

76 

77 mesh_context: Any 

78 

79 def axis_mesh(self, axis: str) -> Any: 

80 """Return a prebuilt one-dimensional mesh without constructing groups. 

81 

82 Args: 

83 axis: cp or ep, in the runtime's established rank order. 

84 """ 

85 if axis == "ep": 

86 region = self.mesh_context.fsdp_moe_mesh 

87 if region is not None: 

88 return region["ep"] 

89 region = self.mesh_context.device_mesh 

90 if region is None or axis not in tuple(region.mesh_dim_names or ()): 

91 raise ValueError(f"no prebuilt {axis!r} mesh is available for parallel binding") 

92 return region[axis] 

93 

94 

95@dataclass(frozen=True) 

96class ParallelBinding: 

97 """One resource object injected through a class-defined bind_parallel method.""" 

98 

99 mesh_binding: MeshBinding 

100 cp_context: CPContext | None = None 

101 ep_context: EPContext | None = None 

102 

103 

104def validate_parallel_module(module: nn.Module, parallel: ParallelModuleSpec, sharding: Any, context: Any) -> None: 

105 """Validate class stages and capabilities after owned-layout port checks. 

106 

107 Args: 

108 module: Actual final member or native component. 

109 parallel: Versioned stage declaration from the class/model spec. 

110 sharding: Final ModuleShardingSpec after user overrides. 

111 context: Immutable logical layout context. 

112 """ 

113 for name, arguments in parallel.stages.items(): 

114 implementation = inspect.getattr_static(type(module), name, None) 

115 if not callable(implementation): 

116 raise ValueError(f"parallel stage {name!r} must be implemented in the module class") 

117 inspect.signature(getattr(module, name)).bind(**dict.fromkeys(arguments)) 

118 if parallel.binding_kind != "none": 

119 implementation = inspect.getattr_static(type(module), "bind_parallel", None) 

120 if not callable(implementation): 

121 raise ValueError("parallel modules must define bind_parallel(binding) in their class") 

122 inspect.signature(module.bind_parallel).bind(None) 

123 if getattr(sharding, "_ep_stack", None): 

124 raise ValueError("explicit expert stages cannot inherit legacy whole-MoE stacking") 

125 if sharding.in_src != sharding.in_dst: 

126 raise ValueError("explicit stage modules must receive their declared computation layout directly") 

127 if sharding.out_dst is not None and sharding.out_src != sharding.out_dst: 

128 raise ValueError("explicit stage modules must restore their declared output layout themselves") 

129 if parallel.binding_kind == "ep": 

130 _validate_expert_ownership(module, parallel, sharding, context) 

131 if parallel.validate is not None: 

132 parallel.validate(module=module, sharding=sharding, context=context) 

133 

134 

135def _validate_expert_ownership(module, parallel, sharding, context): 

136 """Check the bound EP mesh against local and global expert ownership.""" 

137 experts = parallel.options.get("num_experts") 

138 expert_type = type(experts) 

139 if expert_type is not int or experts < 1 or experts % context.axis_sizes["ep"]: 

140 raise ValueError("EP protocols must declare a divisible global num_experts") 

141 sharded = {name: named["ep"] for name, named in sharding.params.items() 

142 if isinstance(named.get("ep"), Shard)} 

143 if not sharded: 

144 raise ValueError("EP transport requires explicit DTensor expert ownership") 

145 for name, placement in sharded.items(): 

146 parameter = module.get_parameter(name) 

147 if placement.dim != 0 or parameter.shape[0] != experts: 

148 raise ValueError("EPContext supports balanced contiguous expert Shard(0) ownership") 

149 

150 

151def bind_parallel_modules(model: nn.Module, plan: Any, mesh_context: Any) -> None: 

152 """Inject shared resources after layout and before activation wrappers/FSDP2. 

153 

154 Args: 

155 model: Final model tree with computation-sharded parameters. 

156 plan: Plan carrying validated family/native protocol declarations. 

157 mesh_context: Existing HyperParallel MeshContext; never a new topology. 

158 """ 

159 mesh_binding = MeshBinding(mesh_context) 

160 prepared = [] 

161 for fqn, parallel in plan.parallel_modules.items(): 

162 if parallel.binding_kind == "none": 

163 continue 

164 module = model.get_submodule(fqn) 

165 existing = getattr(module, "_hp_parallel_binding", None) 

166 if existing is not None: 

167 info = get_family_build_info(module) 

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

169 if (cached is None or not cached.applied or cached.parallel != parallel 

170 or not _same_mesh(existing.mesh_binding.mesh_context, mesh_context)): 

171 raise ValueError(f"module {fqn!r} already has a communication owner") 

172 continue 

173 cp = CPContext(mesh_binding.axis_mesh("cp")) if parallel.binding_kind == "cp" else None 

174 ep = (EPContext(mesh_binding.axis_mesh("ep"), num_experts=parallel.options["num_experts"]) 

175 if parallel.binding_kind == "ep" else None) 

176 prepared.append((module, ParallelBinding(mesh_binding, cp, ep))) 

177 for module, binding in prepared: 

178 module.bind_parallel(binding) 

179 module._hp_parallel_binding = binding # pylint: disable=protected-access 

180 model._hp_parallel_module_specs = dict(plan.parallel_modules) # pylint: disable=protected-access 

181 

182 

183def validate_materialized_modules(model: nn.Module) -> None: 

184 """Check final device/kernel resources after materialization and before compile. 

185 

186 Args: 

187 model: Final loaded model whose class protocols may validate physical resources. 

188 """ 

189 for fqn, parallel in getattr(model, "_hp_parallel_module_specs", {}).items(): 

190 module = model.get_submodule(fqn) 

191 if parallel.validate_materialized is not None: 

192 parallel.validate_materialized(module=module, binding=getattr(module, "_hp_parallel_binding", None))