Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / checkpoint / _weight_check_geometry.py: 96%

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"""Bounded rectangles for checking inverse-converted distributed weights.""" 

16 

17from __future__ import annotations 

18 

19import math 

20from collections.abc import Iterator, Sequence 

21from typing import Any 

22 

23Box = tuple[tuple[int, ...], tuple[int, ...]] 

24 

25 

26def tile_box(offsets: tuple[int, ...], lengths: tuple[int, ...], limit: int) -> Iterator[Box]: 

27 """Split a rectangle into non-overlapping tiles of at most ``limit`` elements.""" 

28 if limit < 1: 

29 raise ValueError("Tile limit must be positive") 

30 if any(size == 0 for size in lengths): 

31 return 

32 if math.prod(lengths) <= limit: 

33 yield offsets, lengths 

34 return 

35 axis = next((axis for axis, size in enumerate(lengths) if size > 1), 0) 

36 step = max(1, limit // math.prod(lengths[axis + 1:])) 

37 for start in range(0, lengths[axis], step): 

38 child_offsets = list(offsets) 

39 child_lengths = list(lengths) 

40 child_offsets[axis] += start 

41 child_lengths[axis] = min(step, lengths[axis] - start) 

42 yield from tile_box(tuple(child_offsets), tuple(child_lengths), limit) 

43 

44 

45def local_block_box(block: Any, shard_offsets: tuple[int, ...], shard_shape: tuple[int, ...]) -> Box | None: 

46 """Intersect an inverse conversion block with one live model shard. 

47 

48 Permutations, splits and stacks map each varying model axis to at most one 

49 output axis. More complex affine maps fail explicitly instead of skipping data. 

50 """ 

51 low = [0] * len(block.lengths) 

52 high = list(block.lengths) 

53 for size, column in zip(block.lengths, block.steps): 

54 if size > 1 and sum(step != 0 for step in column) != 1: 

55 raise ValueError("Inverse conversion has a non-rectangular shard mapping") 

56 for dim, origin in enumerate(block.base): 

57 steps = tuple(column[dim] for column in block.steps) 

58 moving = [(axis, step) for axis, step in enumerate(steps) if step] 

59 start, stop = shard_offsets[dim], shard_offsets[dim] + shard_shape[dim] 

60 if not moving: 

61 if not start <= origin < stop: 

62 return None 

63 continue 

64 if len(moving) != 1 or moving[0][1] != 1: 

65 raise ValueError("Inverse conversion has a non-rectangular shard mapping") 

66 axis, _ = moving[0] 

67 low[axis] = max(low[axis], start - origin) 

68 high[axis] = min(high[axis], stop - origin) 

69 if any(left >= right for left, right in zip(low, high)): 

70 return None 

71 return ( 

72 tuple(origin + left for origin, left in zip(block.offsets, low)), 

73 tuple(right - left for left, right in zip(low, high)), 

74 ) 

75 

76 

77def box_union_volume(boxes: Sequence[Box]) -> int: 

78 """Count unique elements covered by rectangular shards, including replicas.""" 

79 unique = set((tuple(offsets), tuple(lengths)) for offsets, lengths in boxes if all(lengths)) 

80 if not unique: 

81 return 0 

82 if not next(iter(unique))[0]: 

83 return 1 

84 endpoints = sorted({point for offsets, lengths in unique for point in (offsets[0], offsets[0] + lengths[0])}) 

85 return sum( 

86 (right - left) * box_union_volume([ 

87 (offsets[1:], lengths[1:]) for offsets, lengths in unique 

88 if offsets[0] <= left and offsets[0] + lengths[0] >= right 

89 ]) 

90 for left, right in zip(endpoints, endpoints[1:]) 

91 )