Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / ragged_utils.py: 93%
73 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-21 04:29 +0800
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"""RaggedShard geometry adapters for distributed checkpoint planners."""
16from math import prod
17from typing import Any, NamedTuple
19from hyper_parallel.core.distributed_checkpoint.metadata import (
20 ChunkStorageMetadata,
21 MetadataIndex,
22 TensorProperties,
23)
24from hyper_parallel.core.distributed_checkpoint.planner import (
25 WriteItem,
26 WriteItemType,
27)
28from hyper_parallel.core.dtensor._ragged_utils import _compute_ragged_slice
29from hyper_parallel.core.dtensor.dtensor import DTensor
32class RaggedCheckpointBox(NamedTuple):
33 """One regular global box backed by a contiguous local flat interval."""
35 offsets: tuple[int, ...]
36 sizes: tuple[int, ...]
37 local_flat_start: int
38 local_flat_end: int
41def _decompose_flat_interval(
42 shape: tuple[int, ...],
43 flat_start: int,
44 flat_end: int,
45) -> tuple[tuple[tuple[int, ...], tuple[int, ...]], ...]:
46 """Decompose a row-major flat interval into ordered N-D boxes."""
47 total_numel = prod(shape)
48 if not shape:
49 raise ValueError("Ragged checkpoint geometry requires a non-scalar global shape")
50 if flat_start < 0 or flat_end < flat_start or flat_end > total_numel:
51 raise ValueError(
52 "Invalid flat interval for Ragged checkpoint geometry, "
53 f"got interval=({flat_start}, {flat_end}), shape={shape!r}"
54 )
55 if flat_start == flat_end:
56 return ()
58 boxes: list[tuple[tuple[int, ...], tuple[int, ...]]] = []
60 def _decompose_axis(
61 axis: int,
62 start: int,
63 end: int,
64 prefix_offsets: tuple[int, ...],
65 ) -> None:
66 if axis == len(shape) - 1:
67 boxes.append(
68 (
69 prefix_offsets + (start,),
70 (1,) * len(prefix_offsets) + (end - start,),
71 )
72 )
73 return
75 stride = prod(shape[axis + 1:])
76 start_block, start_remainder = divmod(start, stride)
77 end_block = (end - 1) // stride
78 if start_block == end_block:
79 _decompose_axis(
80 axis + 1,
81 start_remainder,
82 start_remainder + end - start,
83 prefix_offsets + (start_block,),
84 )
85 return
87 complete_start = start_block
88 if start_remainder:
89 _decompose_axis(
90 axis + 1,
91 start_remainder,
92 stride,
93 prefix_offsets + (start_block,),
94 )
95 complete_start += 1
97 complete_end, end_remainder = divmod(end, stride)
98 if complete_start < complete_end:
99 boxes.append(
100 (
101 prefix_offsets
102 + (complete_start,)
103 + (0,) * (len(shape) - axis - 1),
104 (1,) * len(prefix_offsets)
105 + (complete_end - complete_start,)
106 + shape[axis + 1:],
107 )
108 )
110 if end_remainder:
111 _decompose_axis(
112 axis + 1,
113 0,
114 end_remainder,
115 prefix_offsets + (complete_end,),
116 )
118 _decompose_axis(0, flat_start, flat_end, ())
119 return tuple(boxes)
122def compute_ragged_boxes(tensor: DTensor) -> tuple[RaggedCheckpointBox, ...]:
123 """Return ordered N-D boxes covering one RaggedShard local flat tensor."""
124 layout = tensor.layout
125 if layout.ragged_shard is None:
126 raise ValueError("compute_ragged_boxes requires a RaggedShard DTensor")
128 global_shape = tuple(tensor.shape)
129 ragged_slice = _compute_ragged_slice(global_shape, layout)
130 raw_boxes = _decompose_flat_interval(
131 global_shape,
132 ragged_slice.flat_start,
133 ragged_slice.flat_end,
134 )
135 boxes: list[RaggedCheckpointBox] = []
136 local_flat_offset = 0
137 for offsets, sizes in raw_boxes:
138 box_numel = prod(sizes)
139 boxes.append(
140 RaggedCheckpointBox(
141 offsets=offsets,
142 sizes=sizes,
143 local_flat_start=local_flat_offset,
144 local_flat_end=local_flat_offset + box_numel,
145 )
146 )
147 local_flat_offset += box_numel
149 if local_flat_offset != ragged_slice.local_numel:
150 raise ValueError(
151 "Ragged checkpoint boxes do not cover the local flat tensor, "
152 f"covered={local_flat_offset}, expected={ragged_slice.local_numel}"
153 )
154 return tuple(boxes)
157def create_ragged_write_items(fqn: str, tensor: DTensor) -> list[WriteItem]:
158 """Create one standard N-D checkpoint write item per RaggedShard box."""
159 local_tensor = tensor.to_local()
160 dtype_str = str(local_tensor.dtype) if hasattr(local_tensor, "dtype") else "unknown"
161 properties = TensorProperties(dtype=dtype_str)
162 items: list[WriteItem] = []
163 for box in compute_ragged_boxes(tensor):
164 chunk = ChunkStorageMetadata(offsets=box.offsets, sizes=box.sizes)
165 items.append(
166 WriteItem(
167 index=MetadataIndex(fqn=fqn, offset=box.offsets, index=None),
168 type=WriteItemType.TENSOR,
169 tensor_data={
170 "chunk": chunk,
171 "properties": properties,
172 "size": tuple(tensor.shape),
173 },
174 )
175 )
176 return items
179def get_ragged_box_tensor(tensor: DTensor, index: MetadataIndex) -> Any:
180 """Return the local flat view corresponding to one global N-D box."""
181 requested_offset = tuple(index.offset) if index.offset is not None else None
182 for box in compute_ragged_boxes(tensor):
183 if box.offsets == requested_offset:
184 local_flat = tensor.to_local().reshape((-1,))
185 return local_flat[
186 box.local_flat_start:box.local_flat_end
187 ].reshape(box.sizes)
188 raise ValueError(
189 "Ragged checkpoint box was not found in the local DTensor, "
190 f"fqn={index.fqn!r}, offset={index.offset!r}"
191 )