Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/core/distributed

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

7# http://www.apache.org/licenses/LICENSE-2.0

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"""Common utility functions.

17What every part of a distributed checkpoint reaches for and no one part owns: the paths,

18the shard geometry, the state dict walk, and the one logger. Checkpoint files are read and

19written in :mod:`checkpoint_io`; replicated shards are sent in :mod:`broadcast`.

20"""

22import time

23from functools import wraps

24from typing import Any, Optional, Union

25from collections.abc import Callable, Collection, Mapping

27import torch

28import torch.distributed as dist

29from torch import Tensor

31from hyper_parallel.core.distributed_checkpoint.metadata import (

32 ChunkStorageMetadata,

33 MetadataIndex,

34 CHUNK_INFO,

35 ChunkInfo,

36)

37from hyper_parallel.core.distributed_checkpoint.planner import SavePlan

38from hyper_parallel.core.distributed_checkpoint.ragged import compute_ragged_boxes

39from hyper_parallel.core.dtensor.layout import infer_slice_area_by_layout

40from hyper_parallel.core.dtensor.dtensor import DTensor

41from hyper_parallel.tools.logging import get_logger

43# The one DCP logger: other distributed_checkpoint modules import this instead of

44# registering a component of their own.

45logger = get_logger("DCP")

48def get_tensor_storage_size(tensor: torch.Tensor) -> int:

49 """

50 The serialized byte size of one tensor.

52 Args:

53 tensor (torch.Tensor): The tensor to measure.

55 Returns:

56 int: Element count times element size, which is what it occupies on disk.

57 """

58 if not isinstance(tensor, torch.Tensor):

59 raise TypeError(f"get_tensor_storage_size expects torch.Tensor, got {type(tensor)!r}")

60 return int(tensor.numel()) * int(tensor.element_size())

63def str_to_dtype(dtype_str: str) -> torch.dtype:

64 """

65 Map a ``torch.<type>`` string from checkpoint metadata back to a ``torch.dtype``.

67 Args:

68 dtype_str (str): The dtype as metadata records it, such as ``torch.float32``.

70 Returns:

71 torch.dtype: The dtype it names.

72 """

73 parts = dtype_str.split(".", 1)

74 if len(parts) != 2:

75 raise ValueError(f"Expected dtype string like 'torch.float32', got {dtype_str!r}.")

76 prefix, name = parts

77 if prefix != "torch":

78 raise ValueError(f"Expected PyTorch dtype string with prefix 'torch', got {dtype_str!r}.")

79 dtype = getattr(torch, name, None)

80 if isinstance(dtype, torch.dtype):

81 return dtype

82 raise ValueError(f"{dtype_str!r} does not resolve to a torch.dtype.")

85def dcp_timer_decorator(func: Callable) -> Callable:

86 """

87 Used to collect statistics on the time consumed in each phase of the DCP.

89 The timings are per-rank, so the rank is part of the message; enable them

90 with ``HP_LOG_CONFIG=DCP:INFO``.

92 Args:

93 func (Callable): The phase to time.

95 Returns:

96 Callable: ``func`` wrapped so that each call logs how long it took.

97 """

98 @wraps(func)

99 def wrapper(*args: Any, **kwargs: Any) -> Any:

100 """Log how long one call of the wrapped phase took, and return its result."""

101 try:

102 rank_id = dist.get_rank()

103 except ValueError:

104 # No process group yet (offline converters, single-process tools).

105 rank_id = 0

106 logger.info("[rank=%d] >>> func %s start exec", rank_id, func.__name__)

107 start_time = time.time()

108 result = func(*args, **kwargs)

109 end_time = time.time()

110 execution_time = end_time - start_time

111 logger.info("[rank=%d] >>> func %s cost %.4f seconds", rank_id, func.__name__, execution_time)

112 return result

113

114 return wrapper

115

116

117def narrow_tensor_by_index(tensor: Any, offsets: tuple, lengths: tuple) -> Any:

118 """

119 Narrow the tensor by (offsets, lengths) per dimension.

120

121 Used for resharding operations to extract a slice from a tensor.

122 Uses plain slice indexing so any tensor-like object works.

123

124 Args:

125 tensor (Any): The tensor to narrow (tensor-like object supporting indexing).

126 offsets (tuple): Tuple of offsets per dimension.

127 lengths (tuple): Tuple of lengths per dimension.

128

129 Returns:

130 Any: The narrowed tensor slice (tensor-like object).

131 """

132 if not offsets or not lengths:

133 return tensor

134 slices = tuple(

135 slice(int(off), int(off) + int(ln))

136 for off, ln in zip(offsets, lengths)

137 )

138 return tensor[slices]

139

140

141def chunk_to_area(chunk: ChunkStorageMetadata) -> tuple[tuple[int, int], ...]:

142 """

143 Convert ChunkStorageMetadata to (start, end) area per dimension.

144

145 Args:

146 chunk (ChunkStorageMetadata): ChunkStorageMetadata instance with offsets and sizes.

147

148 Returns:

149 tuple[tuple[int, int], ...]: Tuple of (start, end) tuples for each dimension.

150 """

151 return tuple(

152 (chunk.offsets[i], chunk.offsets[i] + chunk.sizes[i])

153 for i in range(len(chunk.offsets))

154 )

155

156

157def infer_intersection(

158 area_a: tuple[tuple[int, int], ...],

159 area_b: tuple[tuple[int, int], ...]

160) -> Optional[tuple[tuple[int, int], ...]]:

161 """

162 Calculates the intersection of two tensor slice areas.

163

164 Args:

165 area_a (tuple[tuple[int, int], ...]): First area to intersect.

166 area_b (tuple[tuple[int, int], ...]): Second area to intersect.

167

168 Returns:

169 Optional[tuple[tuple[int, int], ...]]: Tuple of intersection boundaries or None if no intersection.

170 """

171 # Validate input formats

172 def is_valid_axis_list(axis_list: Any) -> None:

173 """Reject an area that is not a sequence of two-element ranges."""

174 if not isinstance(axis_list, (tuple, list)):

175 raise TypeError("Area must be a tuple of ranges")

176 for axis_range in axis_list:

177 if (not isinstance(axis_range, (tuple, list)) \

178 or len(axis_range) != 2):

179 raise TypeError("Each axis range must be a 2-element tuple")

180

181 is_valid_axis_list(area_a)

182 is_valid_axis_list(area_b)

183

184 # Check dimension compatibility

185 if len(area_a) != len(area_b):

186 raise ValueError(

187 f"Area dimension mismatch: {len(area_a)} vs {len(area_b)}"

188 )

189

190 # Calculate intersection for each dimension

191 intersection: list[tuple[int, int]] = []

192 for axis_range_a, axis_range_b in zip(area_a, area_b):

193 left = max(axis_range_a[0], axis_range_b[0])

194 right = min(axis_range_a[1], axis_range_b[1])

195

196 if left >= right: # No intersection in this dimension

197 return None

198

199 intersection.append((left, right))

200

201 return tuple(intersection)

202

203

204def create_chunk_list_for_tensor(obj: Union[Tensor, DTensor]) -> list[ChunkStorageMetadata]:

205 """

206 Create list of local chunks for the given object (DTensor or plain tensor).

207

208 Used to determine what this rank needs to load (resharding).

209

210 Args:

211 obj (Union[Tensor, DTensor]): hyper DTensor or torch Tensor.

212

213 Returns:

214 list[ChunkStorageMetadata]: List of ChunkStorageMetadata representing

215 local chunks needed by this rank.

216 """

217 if isinstance(obj, DTensor):

218 layout = obj.layout

219 if layout is None:

220 shape = obj.shape if hasattr(obj, "shape") else obj.to_local().shape

221 return [ChunkStorageMetadata(offsets=(0,) * len(shape), sizes=tuple(shape))]

222 if layout.ragged_shard is not None:

223 return [

224 ChunkStorageMetadata(offsets=box.offsets, sizes=box.sizes)

225 for box in compute_ragged_boxes(obj)

226 ]

227

228 mesh_shape = getattr(layout, "mesh_shape", None) or getattr(layout, "_mesh", None)

229 tensor_map = getattr(layout, "tensor_map", None) or getattr(layout, "_tensor_map", None)

230 rank_list = getattr(layout, "rank_list", None) or getattr(layout, "_rank_list", None)

231

232 if mesh_shape is None or tensor_map is None or rank_list is None:

233 shape = obj.shape if hasattr(obj, "shape") else obj.to_local().shape

234 return [ChunkStorageMetadata(offsets=(0,) * len(shape), sizes=tuple(shape))]

235

236 current_rank = dist.get_rank()

237 if current_rank not in rank_list:

238 return []

239

240 inner_rank_id = rank_list.index(current_rank)

241 full_shape = obj.shape

242 slice_area = infer_slice_area_by_layout(

243 layout,

244 inner_rank_id,

245 full_shape,

246 )

247 offsets = tuple(s for s, _ in slice_area)

248 sizes = tuple(e - s for s, e in slice_area)

249 return [ChunkStorageMetadata(offsets=offsets, sizes=sizes)]

250

251 if isinstance(obj, Tensor):

252 # handle Tensor with shard information

253 if hasattr(obj, CHUNK_INFO):

254 if not isinstance(getattr(obj, CHUNK_INFO), ChunkInfo):

255 raise ValueError("The attr CHUNK_INFO should be a ChunkInfo instance")

256 chunk = getattr(obj, CHUNK_INFO).chunk

257 return [chunk]

258 # torch.Tensor has exactly one chunk in metadata (full tensor)

259 shape = tuple(obj.shape)

260 return [ChunkStorageMetadata(offsets=(0,) * len(shape), sizes=shape)]

261

262 raise ValueError(f"Not support type {type(obj)} for creating chunk list ")

263

264

265def plan_ownership_masks(

266 all_plans: list[SavePlan],

267 save_to_minimum_rank: bool = False,

268) -> list[bytearray]:

269 """

270 Decide which plan writes each item, as one keep-mask per plan.

271

272 An item present in several plans is redundant: only one plan should write it. Empty

273 uneven shards can share an offset with a nonempty shard; the nonempty shard must win.

274 The owner is the plan with the smallest planned storage so far, or the lowest plan index when

275 ``save_to_minimum_rank`` is True. Ownership is resolved from the plan order and the item

276 metadata and sizes, so every rank running this over the same gathered plans reaches the same answer.

277

278 Duplicates are assigned largest first (longest-processing-time): placing the big shards while

279 the plans are still evenly loaded leaves the small ones to even out the remainder. Assigning

280 in arrival order instead lets a late big shard land on an already-heavy plan, and the

281 checkpoint's wall time is set by whichever plan writes the most.

282

283 Masks are returned instead of filtered plans because the caller walks ``plan.items`` anyway:

284 skipping on a mask costs no hashing and allocates no intermediate copy of every plan.

285

286 Args:

287 all_plans (list[SavePlan]): Local plans gathered from all ranks, indexed by rank.

288 save_to_minimum_rank (bool): If True, assign duplicates to the lowest plan index; else to

289 the plan holding the least data so far. Default False.

290

291 Returns:

292 list[bytearray]: One mask per plan, parallel to that plan's ``items``: 1 marks an item the

293 plan owns and must write, 0 marks a duplicate another plan took.

294 """

295 # index -> [write_item, plan_idx, position, plan_idx, position, ...]. One flat list per

296 # distinct item, so the common unique item costs a single dict lookup rather than an entry

297 # in a duplicate map plus one in a registry plus one in a per-plan set.

298 occurrences_by_index: dict[MetadataIndex, list] = {}

299 for plan_idx, plan in enumerate(all_plans):

300 for position, entry in enumerate(plan.items):

301 occurrences = occurrences_by_index.get(entry.index)

302 entry_size = entry.tensor_storage_size()

303 if entry_size == 0:

304 if occurrences is None:

305 occurrences_by_index[entry.index] = [entry, plan_idx, position]

306 continue

307 if occurrences is None or occurrences[0].tensor_storage_size() == 0:

308 occurrences_by_index[entry.index] = [entry, plan_idx, position]

309 else:

310 if entry.tensor_data != occurrences[0].tensor_data:

311 raise ValueError(f"Checkpoint chunks share {entry.index} with different nonempty metadata")

312 occurrences.append(plan_idx)

313 occurrences.append(position)

314

315 masks = [bytearray(len(plan.items)) for plan in all_plans]

316 storage_sizes = [0] * len(all_plans)

317

318 # Unique items are assigned first so that they are all accounted for in storage_sizes before

319 # duplicates are balanced against those sizes. Sizes are computed here, once per item.

320 duplicates: list[tuple[int, list]] = []

321 for occurrences in occurrences_by_index.values():

322 item_size = occurrences[0].tensor_storage_size() or 1

323 # The layout is [write_item] + (plan_idx, position) * holder_count: drop the slot the

324 # item itself takes, then every two remaining slots are one plan holding it.

325 holder_count = (len(occurrences) - 1) // 2

326 if holder_count > 1:

327 duplicates.append((item_size, occurrences))

328 continue

329 masks[occurrences[1]][occurrences[2]] = 1

330 storage_sizes[occurrences[1]] += item_size

331

332 # Largest first, so the shards with room to unbalance the plans are placed while every plan is

333 # still a candidate. Python's sort is stable, including with reverse=True, so equally sized

334 # duplicates keep their gather order and every rank still agrees on the owner. Sorting is

335 # pointless when every duplicate goes to its lowest plan index regardless.

336 if not save_to_minimum_rank:

337 duplicates.sort(key=lambda pair: pair[0], reverse=True)

338

339 for item_size, occurrences in duplicates:

340 # Occurrences were appended in ascending plan order, so slot 1 is the lowest plan index

341 # and the storage-size search breaks ties towards it.

342 if save_to_minimum_rank:

343 owner_slot = 1

344 else:

345 owner_slot = _least_loaded_owner_slot(occurrences, storage_sizes)

346 owner_idx = occurrences[owner_slot]

347 masks[owner_idx][occurrences[owner_slot + 1]] = 1

348 storage_sizes[owner_idx] += item_size

349

350 return masks

351

352

353def _least_loaded_owner_slot(occurrences: list, storage_sizes: list[int]) -> int:

354 """

355 Pick the occurrence slot whose plan currently holds the least data.

356

357 Args:

358 occurrences (list): ``[write_item, plan_idx, position, ...]`` for one duplicated item.

359 storage_sizes (list[int]): Bytes already assigned to each plan.

360

361 Returns:

362 int: Index into ``occurrences`` of the winning ``plan_idx`` (its position follows it).

363 """

364 # Slot 1 is the first plan_idx and every further holder sits two slots on. Comparing with a

365 # strict < keeps the first minimum, which is the lowest plan index.

366 best_slot = 1

367 best_size = storage_sizes[occurrences[1]]

368 for slot in range(3, len(occurrences), 2):

369 size = storage_sizes[occurrences[slot]]

370 if size < best_size:

371 best_slot, best_size = slot, size

372 return best_slot

373

374

375def traverse_state_dict(

376 state_dict: Any,

377 visitor: Any,

378) -> None:

379 """

380 Invoke ``visitor`` for each value recursively in ``state_dict``.

381 Mapping will be traversed and ``visitor`` will be applied to the leaf elements.

382 ``visitor`` will only be applied to elements in a list or a tuple, if the

383 container contains tensors or mappings.

384 """

385

386 def _is_terminal(value: Any) -> bool:

387 """Leaf-like container: no nested mappings/lists/tuples/tensors to recurse into."""

388 values: Collection

389 if isinstance(value, Mapping):

390 return False

391 if isinstance(value, (list, tuple)):

392 values = value

393 else:

394 return True

395

396 for entry in values:

397 if isinstance(entry, (Mapping, list, tuple)) and not _is_terminal(entry):

398 return False

399 if isinstance(entry, Tensor):

400 return False

401 return True

402

403 def _traverse_obj(path: tuple[Any, ...], value: Any) -> None:

404 if isinstance(value, Mapping):

405 for k, v in value.items():

406 _traverse_obj(path + (str(k),), v)

407 elif _is_terminal(value):

408 visitor(path, value)

409 elif isinstance(value, (list, tuple)):

410 for i, v in enumerate(value):

411 _traverse_obj(path + (i,), v)

412

413 for key, value in state_dict.items():

414 _traverse_obj((str(key),), value)

415

416

417def flatten_state_dict(state_dict: Any) -> tuple[dict[str, Any], dict[str, tuple[Any, ...]]]:

418 """Flatten a nested state dict to dotted FQN keys; returns ``(flat_dict, fqn -> path)``."""

419 fqn_names: dict[str, Any] = {}

420 mappings: dict[str, tuple[Any, ...]] = {}

421

422 def flat_copy(path: tuple[Any, ...], value: Any) -> None:

423 """Record one leaf under its dotted FQN, refusing a name already taken."""

424 new_fqn = ".".join(map(str, path))

425 if new_fqn in fqn_names:

426 raise ValueError(

427 f"Duplicate flattened FQN {new_fqn!r} when converting nested state_dict; "

428 "two different values map to the same dotted name."

429 )

430 fqn_names[new_fqn] = value

431 mappings[new_fqn] = path

432

433 traverse_state_dict(state_dict, flat_copy)

434 return fqn_names, mappings

435

436

437def set_element(root_dict: Any, path: tuple[Any, ...], value: Any) -> None:

438 """Set ``value`` in ``root_dict`` along the ``path`` object path."""

439 if not path:

440 raise ValueError("path must be non-empty")

441 cur_container: Any = root_dict

442

443 def extend_list(lst: list[Any], idx: int) -> None:

444 """Pad ``lst`` with None until ``idx`` is a valid position in it."""

445 while len(lst) <= idx:

446 lst.append(None)

447

448 for i in range(1, len(path)):

449 prev_key = path[i - 1]

450 next_key = path[i]

451 def_val: Any = {} if isinstance(next_key, str) else []

452

453 if isinstance(cur_container, Mapping):

454 cur_container = cur_container.setdefault(prev_key, def_val)

455 else:

456 extend_list(cur_container, prev_key)

457 if cur_container[prev_key] is None:

458 cur_container[prev_key] = def_val

459 cur_container = cur_container[prev_key]

460

461 last_key = path[-1]

462 if isinstance(last_key, int):

463 extend_list(cur_container, last_key)

464

465 cur_container[last_key] = value

466

467

468@dcp_timer_decorator

469def all_gather_object(

470 local_object: Any,

471 world_size: int,

472 use_collectives: bool,

473) -> list[Any]:

474 """

475 Gather objects from all ranks.

476

477 Args:

478 local_object (Any): Local object for current rank.

479 world_size (int): Total number of ranks.

480 use_collectives (bool): Whether to use collective communication.

481

482 Returns:

483 list[Any]: List of all objects from all ranks.

484 """

485 if use_collectives and world_size > 1:

486 all_objects = [None] * world_size

487 dist.all_gather_object(all_objects, local_object)

488 return all_objects

489 return [local_object]

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / utils.py: 92%

210 statements