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

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 

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

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

54 if not named: 

55 return "{}" 

56 items = [ 

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

58 for axis, placement in named.items() 

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

60 ] 

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

62 

63 @staticmethod 

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

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

66 if obj is None: 

67 return "-" 

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

69 if path: 

70 return path 

71 if isinstance(obj, str): 

72 return obj 

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

74 

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

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

77 if spec.params: 

78 lines.append(" parameter sharding:") 

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

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

81 else: 

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

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

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

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

86 if attr_plan.auto_divide: 

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

88 if attr_plan.user_divide: 

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

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

91 lines.append( 

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

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

94 ) 

95 

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

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

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

99 # in _builder since changeset 4c. 

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

101 PrecompiledBoundary, 

102 ) 

103 

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

105 if boundary.in_plan: 

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

107 for op in boundary.in_plan: 

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

109 lines.append( 

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

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

112 ) 

113 if boundary.out_plan: 

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

115 for op in boundary.out_plan: 

116 lines.append( 

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

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

119 ) 

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

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

122 

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

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

125 injection = [] 

126 if spec.local_compute_fn is not None: 

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

128 if spec.inner_wrapper is not None: 

129 injection.append( 

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

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

132 ) 

133 if spec.inner_out_src is not None: 

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

135 if spec.region_dispatch is not None: 

136 meaning = ( 

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

138 if spec.region_dispatch is False 

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

140 ) 

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

142 if injection: 

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

144 else: 

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

146 

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

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

149 handlers = { 

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

151 } 

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

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

154 

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

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

157 

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

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

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

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

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

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

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

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

166 docs up front. 

167 

168 Args: 

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

170 None reports all boundaries. 

171 """ 

172 lines = [ 

173 "=== ShardingPlan introspection report ===", 

174 f"mesh_dim_names={self.mesh_dim_names} " 

175 f"sequence_parallel={self.sequence_parallel} " 

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

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

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

179 ] 

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

181 lines.append( 

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

183 "(existing boundaries listed above)") 

184 return "\n".join(lines) 

185 selected = ( 

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

187 

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

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

190 self._append_parameter_explanation(lines, spec) 

191 self._append_communication_explanation(lines, spec) 

192 self._append_injection_explanation(lines, spec) 

193 self._append_handler_explanation(lines, name) 

194 return "\n".join(lines)