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

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# pylint: disable=missing-apache-license-header 

16"""plan: the model-level sharding plan (05 §3.1 canonical). 

17 

18:class:`ShardingPlan` is the single authoritative model-level plan of 

19AutoModels — it is NOT a port of the legacy ``core/shard/sharding_plan.py`` 

20(the two data models are incompatible and must never import each other). 

21""" 

22 

23from dataclasses import dataclass, field 

24from typing import Any, Dict, List, Optional, Tuple 

25 

26from hyper_parallel.distributed.recipe_spec import ( 

27 ModuleShardingSpec, 

28 NamedPlacement, 

29) 

30 

31 

32@dataclass 

33class ShardingPlan: 

34 """Complete sharding plan for a model (05 §3.1).""" 

35 # Contains both activation boundaries and parameter-only source specs. 

36 modules: Dict[str, ModuleShardingSpec] = field(default_factory=dict) 

37 

38 # Global switches 

39 sequence_parallel: bool = True 

40 loss_parallel: bool = False 

41 

42 # Special parameter handlers: {module_fqn.param_name: handler_name} 

43 special_handlers: Dict[str, str] = field(default_factory=dict) 

44 

45 # Mesh dimension names (consistent with DeviceMesh.mesh_dim_names) 

46 mesh_dim_names: Tuple[str, ...] = () 

47 

48 # Tied-weight pairs: [(fqn_a, fqn_b)], parameters sharing storage 

49 # (embed_tokens <-> lm_head). 

50 tied_pairs: List[Tuple[str, str]] = field(default_factory=list) 

51 parallel_modules: Dict[str, Any] = field(default_factory=dict) 

52 owned_parameter_names: set[str] = field(default_factory=set) 

53 parameter_aliases: List[Tuple[str, str]] = field(default_factory=list) 

54 

55 def _format_named_placement(self, named: Optional[NamedPlacement]) -> str: 

56 """Format a named placement as a compact one-line dictionary.""" 

57 if not named: 

58 return "{}" 

59 items = [ 

60 (getattr(axis, "value", axis), placement) 

61 for axis, placement in named.items() 

62 if not self.mesh_dim_names or getattr(axis, "value", axis) in self.mesh_dim_names 

63 ] 

64 return "{" + ", ".join(f"{axis}: {placement!r}" for axis, placement in items) + "}" 

65 

66 @staticmethod 

67 def _format_injection_callable(obj: Any) -> str: 

68 """Format a callable, Target, or registry name for display.""" 

69 if obj is None: 

70 return "-" 

71 path = getattr(obj, "_target_path", None) 

72 if path: 

73 return path 

74 if isinstance(obj, str): 

75 return obj 

76 return getattr(obj, "__qualname__", repr(obj)) 

77 

78 def _append_parameter_explanation(self, lines: List[str], spec: ModuleShardingSpec) -> None: 

79 """Append parameter, TP-local attribute, and deferred-bias details.""" 

80 if spec.params: 

81 lines.append(" parameter sharding:") 

82 for param_name, named in spec.params.items(): 

83 lines.append(f" {param_name}: {self._format_named_placement(named)}") 

84 else: 

85 lines.append(" parameter sharding: none ({} = this boundary shards no parameters, I/O stitching only)") 

86 attr_plan = spec._tp_local_attr_plan # pylint: disable=protected-access 

87 if attr_plan is not None and (attr_plan.auto_divide or attr_plan.user_divide): 

88 lines.append(" TP-local attribute division:") 

89 if attr_plan.auto_divide: 

90 lines.append(" auto(D-17): " + ", ".join(attr_plan.auto_divide)) 

91 if attr_plan.user_divide: 

92 lines.append(" user(plan_overrides): " + ", ".join(attr_plan.user_divide)) 

93 if spec._deferred_bias_params: # pylint: disable=protected-access 

94 lines.append( 

95 " deferred bias (D-22, no bias inside the region, added exactly once after the TP reduction): " 

96 + ", ".join(spec._deferred_bias_params) # pylint: disable=protected-access 

97 ) 

98 

99 def _append_communication_explanation(self, lines: List[str], spec: ModuleShardingSpec) -> None: 

100 """Append compiled input and output redistribution operations.""" 

101 # Lazy import keeps plan.py import-light; PrecompiledBoundary lives 

102 # in _builder since changeset 4c. 

103 from hyper_parallel.distributed._builder.precompiled_boundary import ( # pylint: disable=C0415 

104 PrecompiledBoundary, 

105 ) 

106 

107 boundary = PrecompiledBoundary(spec, None, self.mesh_dim_names) 

108 if boundary.in_plan: 

109 lines.append(" input communication plan (in_src -> in_dst):") 

110 for op in boundary.in_plan: 

111 tag = "passthrough" if op.collective_type == "identity" else op.collective_type 

112 lines.append( 

113 f" {op.arg_name}: {tuple(map(repr, op.src_placements))} -> " 

114 f"{tuple(map(repr, op.dst_placements))} [{tag}]" 

115 ) 

116 if boundary.out_plan: 

117 lines.append(" output communication plan (out_src -> out_dst):") 

118 for op in boundary.out_plan: 

119 lines.append( 

120 f" {op.arg_name}(tuple[{op.arg_index}]): {tuple(map(repr, op.src_placements))} -> " 

121 f"{tuple(map(repr, op.dst_placements))} [{op.collective_type}]" 

122 ) 

123 if not boundary.in_plan and not boundary.out_plan: 

124 lines.append(" boundary communication: none") 

125 

126 def _append_injection_explanation(self, lines: List[str], spec: ModuleShardingSpec) -> None: 

127 """Append compute-injection declarations and dispatch semantics.""" 

128 injection = [] 

129 if spec.local_compute_fn is not None: 

130 injection.append(f"local_compute_fn={self._format_injection_callable(spec.local_compute_fn)}") 

131 if spec.inner_wrapper is not None: 

132 injection.append( 

133 f"inner_wrapper={self._format_injection_callable(spec.inner_wrapper)}" 

134 f"(target={spec.inner_target or 'auto'})" 

135 ) 

136 if spec.inner_out_src is not None: 

137 injection.append(f"inner_out_src={spec.inner_out_src}") 

138 if spec.region_dispatch is not None: 

139 meaning = ( 

140 "black-box managed (propagation check skipped inside the region, declarative re-wrap)" 

141 if spec.region_dispatch is False 

142 else "dispatch-through (real validation under validate is enabled)" 

143 ) 

144 injection.append(f"region_dispatch={spec.region_dispatch} -> {meaning}") 

145 if injection: 

146 lines.append(" injection: " + "; ".join(injection)) 

147 else: 

148 lines.append(" injection: none (ordinary boundary, dispatch-through under validate)") 

149 

150 def _append_handler_explanation(self, lines: List[str], boundary_fqn: str) -> None: 

151 """Append special parameter handlers belonging to one boundary.""" 

152 handlers = { 

153 key: value for key, value in self.special_handlers.items() if key.startswith(boundary_fqn + ".") 

154 } 

155 for key, handler in handlers.items(): 

156 lines.append(f" special handling: {key[len(boundary_fqn) + 1:]} -> {handler}") 

157 

158 def explain(self, fqn: Optional[str] = None) -> str: 

159 """Human-readable introspection report of this plan (usability tool). 

160 

161 Per boundary: the parameter sharding table (param → placement), the 

162 compiled boundary communication plan (in/out RedistOps — which 

163 tensor, from which layout to which, which collective), the injection 

164 declarations and their resolution result, and any special-handler 

165 entries. All of this information already lives in the plan/specs — 

166 this is purely a formatting outlet. The intended learning path: 

167 read your own model's actual sharding first, then reverse-engineer 

168 the concepts — instead of building the layout mental model from 

169 docs up front. 

170 

171 Args: 

172 fqn: optional exact boundary FQN — report just that boundary; 

173 None reports all boundaries. 

174 """ 

175 lines = [ 

176 "=== ShardingPlan introspection report ===", 

177 f"mesh_dim_names={self.mesh_dim_names} " 

178 f"sequence_parallel={self.sequence_parallel} " 

179 f"loss_parallel={self.loss_parallel}", 

180 f"boundaries: {len(self.modules)} tied_pairs: " 

181 + (", ".join(f"{a}<->{b}" for a, b in self.tied_pairs) or "none"), 

182 ] 

183 if fqn is not None and fqn not in self.modules: 

184 lines.append( 

185 f"\n[!] {fqn!r} is not a boundary of this plan " 

186 "(existing boundaries listed above)") 

187 return "\n".join(lines) 

188 selected = ( 

189 {fqn: self.modules[fqn]} if fqn is not None else self.modules) 

190 

191 for name, spec in selected.items(): 

192 lines.append(f"\n[{name}]") 

193 self._append_parameter_explanation(lines, spec) 

194 self._append_communication_explanation(lines, spec) 

195 self._append_injection_explanation(lines, spec) 

196 self._append_handler_explanation(lines, name) 

197 return "\n".join(lines)