Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / dmodule / sharding.py: 100%

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"""Declarative sharding specs and placement resolution for dmodule. 

16 

17Provides :class:`ShardingConfig`, :class:`LocalMapConfig`, and 

18:func:`resolve_placements` to convert :data:`~hyper_parallel.dmodule.types.NamedPlacement` 

19maps into mesh-axis-ordered placement lists for 

20:meth:`~hyper_parallel.dmodule.module.Module.parallelize`. 

21 

22Example: 

23 Column-parallel weight with replicated activations:: 

24 

25 from hyper_parallel.core.dtensor.placement_types import Replicate, Shard 

26 from hyper_parallel.dmodule.sharding import ShardingConfig 

27 from hyper_parallel.dmodule.types import MeshAxisName 

28 

29 ShardingConfig( 

30 state_shardings={"weight": {MeshAxisName.TP: Shard(0)}}, 

31 in_dst_shardings={"x": {MeshAxisName.TP: Replicate()}}, 

32 ) 

33""" 

34 

35from dataclasses import dataclass, field 

36 

37from hyper_parallel.core.dtensor.placement_types import Placement, Replicate 

38from hyper_parallel.dmodule.types import MeshAxisName, NamedPlacement 

39 

40__all__ = [ 

41 "LocalMapConfig", 

42 "NamedPlacement", 

43 "ShardingConfig", 

44 "resolve_placements", 

45] 

46 

47 

48@dataclass(kw_only=True, slots=True) 

49class LocalMapConfig: 

50 """Gradient placement specs for ``local_map`` modules (wired in a future milestone). 

51 

52 Args: 

53 in_grad_placements: One :class:`~hyper_parallel.dmodule.types.NamedPlacement` 

54 per forward input that uses ``local_map``. 

55 

56 Example:: 

57 

58 from hyper_parallel.core.dtensor.placement_types import Replicate, Shard 

59 from hyper_parallel.dmodule.sharding import LocalMapConfig, ShardingConfig 

60 from hyper_parallel.dmodule.types import MeshAxisName 

61 

62 sharding = ShardingConfig( 

63 state_shardings={"weight": {MeshAxisName.TP: Shard(0)}}, 

64 local_map=LocalMapConfig( 

65 in_grad_placements=({MeshAxisName.TP: Replicate()},), 

66 ), 

67 ) 

68 """ 

69 

70 in_grad_placements: tuple[NamedPlacement, ...] 

71 

72 def to_dict(self) -> dict: 

73 """Serialize to a JSON-friendly dict (debug/logging). 

74 

75 Returns: 

76 Dict with a ``repr`` field for the whole spec. 

77 """ 

78 return {"repr": repr(self)} 

79 

80 

81@dataclass(kw_only=True, slots=True) 

82class ShardingConfig: 

83 """Declarative sharding for one :class:`~hyper_parallel.dmodule.module.Module`. 

84 

85 Field names for ``in_*`` / ``out_*`` must match ``forward`` argument names 

86 (positional args are mapped by name inside :meth:`Module.parallelize`). 

87 

88 Args: 

89 state_shardings: Parameter path (e.g. ``"weight"``) → :data:`NamedPlacement`. 

90 in_src_shardings: Optional source layout for each forward input before adapt. 

91 in_dst_shardings: Target layout for each forward input (redistribute if needed). 

92 out_src_shardings: Optional source layout(s) for forward output(s). 

93 out_dst_shardings: Target layout for a single tensor output. 

94 local_map: Optional :class:`LocalMapConfig` (not yet applied in ``parallelize``). 

95 

96 Example: 

97 TP column-parallel weight with replicated activations:: 

98 

99 from hyper_parallel.core.dtensor.placement_types import Replicate, Shard 

100 from hyper_parallel.dmodule.sharding import ShardingConfig 

101 from hyper_parallel.dmodule.types import MeshAxisName 

102 

103 ShardingConfig( 

104 state_shardings={ 

105 "weight": {MeshAxisName.TP: Shard(0)}, 

106 }, 

107 in_dst_shardings={ 

108 "x": {MeshAxisName.TP: Replicate()}, 

109 }, 

110 out_dst_shardings={MeshAxisName.TP: Shard(-1)}, 

111 ) 

112 """ 

113 

114 state_shardings: dict[str, NamedPlacement] = field(default_factory=dict) 

115 in_src_shardings: dict[str, NamedPlacement] | None = None 

116 in_dst_shardings: dict[str, NamedPlacement] | None = None 

117 out_src_shardings: NamedPlacement | tuple[NamedPlacement, ...] | None = None 

118 out_dst_shardings: NamedPlacement | None = None 

119 local_map: LocalMapConfig | None = None 

120 

121 def to_dict(self) -> dict: 

122 """Serialize to a JSON-friendly dict (debug/logging). 

123 

124 Returns: 

125 Dict with a ``repr`` field for the whole spec. 

126 """ 

127 return {"repr": repr(self)} 

128 

129 

130def _axis_key(axis: MeshAxisName | str) -> str: 

131 """Normalize a mesh axis key to a plain string for lookup.""" 

132 if isinstance(axis, MeshAxisName): 

133 return axis.value 

134 return str(axis) 

135 

136 

137def resolve_placements( 

138 named: NamedPlacement, 

139 mesh_axis_names: tuple[str, ...], 

140) -> list[Placement]: 

141 """Convert :data:`NamedPlacement` to an ordered placement list for DTensor APIs. 

142 

143 The returned list aligns with ``mesh_axis_names`` (same order and length). 

144 Axes listed on the mesh but omitted from *named* default to 

145 :class:`~hyper_parallel.core.dtensor.placement_types.Replicate`. Keys in *named* 

146 that are not present on the mesh (including case mismatches) raise 

147 :class:`ValueError`. 

148 

149 Args: 

150 named: Placement per mesh axis name (:class:`MeshAxisName` or lowercase str). 

151 mesh_axis_names: Names from ``DeviceMesh.mesh_dim_names`` in axis order. 

152 

153 Returns: 

154 One :class:`~hyper_parallel.core.dtensor.placement_types.Placement` per mesh 

155 axis, suitable for :func:`~hyper_parallel.core.dtensor.dtensor.distribute_tensor` 

156 and ``DTensor.redistribute``. 

157 

158 Raises: 

159 ValueError: If *named* declares an axis that does not appear on the mesh. 

160 

161 Example: 

162 Partial TP spec on a ``(dp, tp)`` mesh — ``dp`` is filled with ``Replicate``:: 

163 

164 from hyper_parallel.core.dtensor.placement_types import Replicate, Shard 

165 from hyper_parallel.dmodule.sharding import resolve_placements 

166 from hyper_parallel.dmodule.types import MeshAxisName 

167 

168 mesh_axis_names = ("dp", "tp") 

169 named = {MeshAxisName.TP: Shard(0)} 

170 

171 placements = resolve_placements(named, mesh_axis_names) 

172 assert placements[0].is_replicate() 

173 assert placements[1].is_shard(0) 

174 

175 Case-sensitive matching — ``"TP"`` does not match ``"tp"``:: 

176 

177 resolve_placements({"TP": Shard(0)}, ("tp",)) # raises ValueError 

178 """ 

179 normalized: dict[str, Placement] = { 

180 _axis_key(axis): placement for axis, placement in named.items() 

181 } 

182 result: list[Placement] = [] 

183 for axis_name in mesh_axis_names: 

184 if axis_name in normalized: 

185 result.append(normalized[axis_name]) 

186 else: 

187 result.append(Replicate()) 

188 

189 declared = set(normalized.keys()) 

190 extra = declared - set(mesh_axis_names) 

191 if extra: 

192 raise ValueError(f"NamedPlacement has axes not in mesh: {extra}") 

193 return result