Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / reshard.py: 74%
136 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-25 04:27 +0800
1# Copyright 2025 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"""resharding tensor"""
16import operator
17from typing import Any, Optional
18from functools import reduce
19import numpy as np
21from hyper_parallel.core.dtensor.layout import (
22 Layout,
23 infer_slice_area_by_rank,
24)
27def check_layout(layout: Optional[Any], name: str) -> None:
28 """
29 Validates that a layout contains required attributes with correct types.
31 Args:
32 layout (Optional[Any]): Layout object to validate.
33 name (str): Name of the layout (for error messages).
35 Raises:
36 ValueError: If layout missing required attributes or has size mismatches
37 TypeError: If layout components are not tuples/lists
38 """
39 if not layout:
40 return
42 # Check for required attributes
43 required_attrs = ['mesh_shape', '_tensor_map', '_rank_list']
44 for attr in required_attrs:
45 if not hasattr(layout, attr):
46 raise ValueError(
47 f"Layout {name} must contain attribute {attr}"
48 )
50 # Validate component types
51 def check_type_is_sequence(obj: Any, obj_name: str) -> None:
52 if not isinstance(obj, (tuple, list)):
53 raise TypeError(
54 f"Layout {name} {obj_name} must be tuple or list, "
55 f"but got {type(obj).__name__}"
56 )
58 layout_dict = layout.to_dict()
59 check_type_is_sequence(layout_dict['mesh_shape'], 'mesh_shape')
60 check_type_is_sequence(layout_dict['tensor_map'], 'tensor_map')
61 check_type_is_sequence(layout_dict['rank_list'], 'rank_list')
63 # Validate rank list size matches device count
64 dev_num = reduce(operator.mul, layout_dict['mesh_shape'])
65 if len(layout_dict['rank_list']) != dev_num:
66 raise ValueError(
67 f"Layout {name} rank_list size ({len(layout_dict['rank_list'])}) "
68 f"must match device count ({dev_num})"
69 )
72def rank_id_to_dev_id_list(mesh_shape: tuple[int, ...], rank_id: int) -> list[int]:
73 """
74 Converts a rank ID to a list of device IDs based on the mesh shape.
76 Args:
77 mesh_shape (tuple[int, ...]): Shape of the mesh shape.
78 rank_id (int): Global rank ID to convert.
80 Returns:
81 list[int]: List of device IDs corresponding to the rank.
82 """
83 dims = len(mesh_shape)
84 dev_id_list = [0] * dims
86 for i in range(dims - 1, -1, -1):
87 dev_id_list[i] = rank_id % mesh_shape[i]
88 rank_id = rank_id // mesh_shape[i]
90 return dev_id_list
93def infer_intersection(
94 area_a: tuple[tuple[int, int], ...],
95 area_b: tuple[tuple[int, int], ...]
96) -> Optional[tuple[tuple[int, int], ...]]:
97 """
98 Calculates the intersection of two tensor slice areas.
100 Args:
101 area_a (tuple[tuple[int, int], ...]): First area to intersect.
102 area_b (tuple[tuple[int, int], ...]): Second area to intersect.
104 Returns:
105 Optional[tuple[tuple[int, int], ...]]: Tuple of intersection boundaries or None if no intersection.
106 """
107 # Validate input formats
108 def is_valid_axis_list(axis_list: Any) -> None:
109 if not isinstance(axis_list, (tuple, list)):
110 raise TypeError("Area must be a tuple of ranges")
111 for axis_range in axis_list:
112 if (not isinstance(axis_range, (tuple, list)) \
113 or len(axis_range) != 2):
114 raise TypeError("Each axis range must be a 2-element tuple")
116 is_valid_axis_list(area_a)
117 is_valid_axis_list(area_b)
119 # Check dimension compatibility
120 if len(area_a) != len(area_b):
121 raise ValueError(
122 f"Area dimension mismatch: {len(area_a)} vs {len(area_b)}"
123 )
125 # Calculate intersection for each dimension
126 intersection: list[tuple[int, int]] = []
127 for axis_range_a, axis_range_b in zip(area_a, area_b):
128 left = max(axis_range_a[0], axis_range_b[0])
129 right = min(axis_range_a[1], axis_range_b[1])
131 if left >= right: # No intersection in this dimension
132 return None
134 intersection.append((left, right))
136 return tuple(intersection)
139class ReshardHandler:
140 """
141 Handles tensor resharding between different distributed layouts.
143 This class manages the process of reshaping and redistributing tensors between
144 different parallel layouts. It calculates necessary tensor slices, validates
145 input layouts, and assembles the final tensor for the target rank.
147 Args:
148 param_name (str): Name of the parameter (without pipeline stage prefix).
149 full_shape (tuple[int, ...]): Complete shape of the tensor before sharding.
150 from_layout (Optional[Any]): Source layout containing mesh shape, tensor map, and rank list.
151 to_layout (Optional[Any]): Target layout containing mesh shape, tensor map, and rank list.
152 to_rank_id (int): Target rank ID to receive the resharded tensor.
154 Raises:
155 ValueError: If both layouts are None or layouts contain invalid attributes
156 TypeError: If layout components are not tuples/lists
157 """
158 def __init__(
159 self,
160 param_name: str,
161 full_shape: tuple[int, ...],
162 from_layout: Optional[Any],
163 to_layout: Optional[Any],
164 to_rank_id: int
165 ):
166 # Validate input layouts
167 check_layout(from_layout, 'from_layout')
168 check_layout(to_layout, 'to_layout')
170 if from_layout is None and to_layout is None:
171 raise ValueError("`from_layout` and `to_layout` cannot both be None.")
173 # Initialize basic attributes
174 self.param_name = param_name
175 self.full_shape = full_shape
176 self.from_uneven_shard_mesh_dims = (
177 from_layout.uneven_shard_mesh_dims
178 if isinstance(from_layout, Layout)
179 else ()
180 )
181 self.to_uneven_shard_mesh_dims = (
182 to_layout.uneven_shard_mesh_dims
183 if isinstance(to_layout, Layout)
184 else ()
185 )
187 # Process source layout configuration
188 if from_layout is None:
189 self.from_mesh_shape = (1,)
190 self.from_tensor_map = tuple(0 for _ in full_shape)
191 self.from_rank_list = [0]
192 else:
193 from_layout_dict = from_layout.to_dict()
194 self.from_mesh_shape = from_layout_dict["mesh_shape"]
195 self.from_tensor_map = from_layout_dict["tensor_map"]
196 self.from_rank_list = from_layout_dict["rank_list"]
198 # Process target layout configuration
199 if to_layout is None:
200 self.to_mesh_shape = (1,)
201 self.to_tensor_map = tuple(0 for _ in full_shape)
202 self.to_rank_list = [0]
203 self.to_rank_id = 0
204 else:
205 to_layout_dict = to_layout.to_dict()
206 self.to_mesh_shape = to_layout_dict["mesh_shape"]
207 self.to_tensor_map = to_layout_dict["tensor_map"]
208 self.to_rank_list = to_layout_dict["rank_list"]
209 self.to_rank_id = to_rank_id
210 if self.to_rank_id not in self.to_rank_list:
211 raise ValueError("Input to_rank_id is not in to_rank_list.")
213 # Calculate device counts and internal rank mappings
214 self.from_dev_num = len(self.from_rank_list)
215 self.inner_from_rank_list = range(self.from_dev_num)
216 self.inner_to_rank_id = self.to_rank_list.index(self.to_rank_id)
218 # Compute redundancy information
219 self.inner_deredundancy_from_rank_list = (
220 self._infer_inner_deredundancy_rank_list_by_from_layout()
221 if from_layout else [0]
222 )
223 self.global_union_area_map: dict[int, tuple[tuple[int, int], ...]] = {}
224 self.to_area = () # Initialized in infer_all_tensor_offset()
226 def _infer_inner_deredundancy_rank_list_by_from_layout(self) -> list[int]:
227 """
228 Infers ranks containing non-redundant data from the source layout.
230 Returns:
231 List of ranks with unique data slices
232 """
233 inner_deredundancy_rank_list: list[int] = []
234 dev_dim = len(self.from_mesh_shape)
236 # Collect relevant device dimensions from tensor map
237 from_dev_map = set()
238 for map_dev in self.from_tensor_map:
239 if isinstance(map_dev, (list, tuple)):
240 for map_dev_inner in map_dev:
241 from_dev_map.add(dev_dim - map_dev_inner - 1)
242 else:
243 from_dev_map.add(dev_dim - map_dev - 1)
245 # Filter ranks with non-redundant data
246 unused_dims = [dim for dim in range(dev_dim) if dim not in from_dev_map]
247 if not unused_dims:
248 return list(self.inner_from_rank_list)
249 for rank_id in self.inner_from_rank_list:
250 dev_id_list = rank_id_to_dev_id_list(self.from_mesh_shape, rank_id)
251 # check redundant
252 found_redundant = False
253 for dim in unused_dims:
254 if dev_id_list[dim] > 0:
255 found_redundant = True
256 break
258 # save not redundant rank
259 if not found_redundant:
260 inner_deredundancy_rank_list.append(rank_id)
262 return inner_deredundancy_rank_list
264 def infer_all_tensor_offset(self) -> dict[int, tuple[tuple[int, int], ...]]:
265 """
266 Calculates required tensor slices from each source rank.
268 Determines which parts of the tensor need to be collected from each source
269 rank to assemble the target tensor slice.
271 Returns:
272 Dictionary mapping source ranks to their required slice offsets
273 """
274 # Calculate target area for current rank
275 self.to_area = infer_slice_area_by_rank(
276 self.to_mesh_shape,
277 self.to_tensor_map,
278 self.inner_to_rank_id,
279 self.full_shape,
280 self.to_uneven_shard_mesh_dims,
281 )
283 # Calculate required slices from each source rank
284 local_union_areas_map: dict[int, tuple[tuple[int, int], ...]] = {}
285 self.global_union_area_map.clear()
287 for inner_rank_id in self.inner_deredundancy_from_rank_list:
288 # Get source area for this rank
289 from_area = infer_slice_area_by_rank(
290 self.from_mesh_shape,
291 self.from_tensor_map,
292 inner_rank_id,
293 self.full_shape,
294 self.from_uneven_shard_mesh_dims,
295 )
297 # Find overlapping area between source and target
298 union_area = infer_intersection(from_area, self.to_area)
299 if union_area is not None:
300 source_rank = self.from_rank_list[inner_rank_id]
301 self.global_union_area_map[source_rank] = union_area
303 # Calculate relative offsets within source slice
304 local_union_areas_map[source_rank] = tuple(
305 (union_range[0] - from_range[0], union_range[1] - from_range[0])
306 for union_range, from_range in zip(union_area, from_area)
307 )
309 return local_union_areas_map
311 def get_real_tensor(self, from_tensor_map: dict[int, np.ndarray]) -> np.ndarray:
312 """
313 Assembles the final tensor for the target rank from collected slices.
315 Args:
316 from_tensor_map (dict[int, np.ndarray]): Dictionary mapping source ranks to their tensor slices.
318 Returns:
319 np.ndarray: Assembled tensor for the target rank.
321 Raises:
322 ValueError: If input slices are missing or have incorrect shapes
323 """
324 if not from_tensor_map:
325 raise ValueError("Input from_tensor_map cannot be empty")
327 # Validate input slices
328 for from_rank_id, from_area in self.global_union_area_map.items():
329 if from_rank_id not in from_tensor_map:
330 raise ValueError(
331 f"Missing slice data from rank {from_rank_id}. "
332 "Please provide all required slices from infer_all_tensor_offset."
333 )
335 # Validate slice shape matches expected size
336 expected_shape = tuple(end - start for start, end in from_area)
337 actual_shape = from_tensor_map[from_rank_id].shape
338 if expected_shape != actual_shape:
339 raise ValueError(
340 f"Slice from rank {from_rank_id} has incorrect shape. "
341 f"Expected {expected_shape}, got {actual_shape}."
342 )
344 # Create target tensor and assign slices
345 to_slice_shape = [end - start for start, end in self.to_area]
346 dtype = next(iter(from_tensor_map.values())).dtype
347 real_tensor = np.zeros(to_slice_shape, dtype=dtype)
349 for from_rank_id, from_slice in from_tensor_map.items():
350 from_area = self.global_union_area_map[from_rank_id]
352 # Calculate assignment indices in target tensor
353 assign_slices = tuple(
354 slice(from_axis[0] - to_axis[0], from_axis[1] - to_axis[0])
355 for from_axis, to_axis in zip(from_area, self.to_area)
356 )
358 real_tensor[assign_slices] = from_slice
360 return real_tensor