Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed_data / model_batch.py: 94%

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"""Locally declared tensor layouts for serialization-free model broadcasts.""" 

16 

17from __future__ import annotations 

18 

19import copy 

20from collections.abc import Mapping 

21from dataclasses import dataclass 

22from typing import Any 

23 

24import torch # pylint: disable=forbidden-backend-import 

25 

26 

27@dataclass(frozen=True) 

28class TensorFieldSpec: 

29 """Describe one tensor field identically on every model-parallel rank. 

30 

31 Args: 

32 shape: Fixed dimensions, with None for dimensions sent in a small 

33 integer header on each batch. An empty tuple denotes a scalar. 

34 dtype: Required tensor dtype; broadcasts never convert its values. 

35 device: "device" for the rank-local training device, or "cpu" for 

36 fields that must remain on Host after broadcast. 

37 """ 

38 

39 shape: tuple[int | None, ...] 

40 dtype: torch.dtype 

41 device: str = "device" 

42 

43 def __post_init__(self) -> None: 

44 if not isinstance(self.shape, tuple) or any( 

45 dim is not None and (isinstance(dim, bool) or not isinstance(dim, int) or dim < 0) 

46 for dim in self.shape 

47 ): 

48 raise ValueError("TensorFieldSpec.shape must be a tuple of non-negative integers or None.") 

49 if not isinstance(self.dtype, torch.dtype): 

50 raise ValueError("TensorFieldSpec.dtype must be a torch.dtype.") 

51 if self.device not in ("device", "cpu"): 

52 raise ValueError("TensorFieldSpec.device must be 'device' or 'cpu'.") 

53 

54 

55class ModelBatchLayout: 

56 """Validate and rebuild a fixed local tree without transmitting its structure.""" 

57 

58 def __init__(self, spec: Any) -> None: 

59 """Freeze the local declaration and enumerate its tensor fields. 

60 

61 Args: 

62 spec: Nested tensor field declarations and static scalar values. 

63 """ 

64 self.spec = copy.deepcopy(spec) 

65 self.fields = tuple(self._fields(self.spec)) 

66 self.dynamic_dimensions = sum(dim is None for field in self.fields for dim in field.shape) 

67 

68 @classmethod 

69 def _fields(cls, spec: Any) -> Any: 

70 """Validate nested layout leaves and yield tensor specifications in traversal order.""" 

71 if isinstance(spec, TensorFieldSpec): 

72 yield spec 

73 elif isinstance(spec, Mapping): 

74 if any(not isinstance(key, str) for key in spec): 

75 raise ValueError("model_batch_spec mapping keys must be strings.") 

76 for item in spec.values(): 

77 yield from cls._fields(item) 

78 elif isinstance(spec, (tuple, list)): 

79 for item in spec: 

80 yield from cls._fields(item) 

81 elif spec is not None and not isinstance(spec, (bool, int, float, str)): 

82 raise ValueError("model_batch_spec leaves must be TensorFieldSpec or static scalar values.") 

83 

84 def flatten(self, batch: Any) -> list[torch.Tensor]: 

85 """Validate source fields and return tensors in the declared order. 

86 

87 Args: 

88 batch: Source batch matching this layout. 

89 """ 

90 return list(self._flatten(self.spec, batch)) 

91 

92 @classmethod 

93 def _flatten(cls, spec: Any, value: Any) -> Any: 

94 """Validate batch structure and static values while yielding its tensor leaves.""" 

95 if isinstance(spec, TensorFieldSpec): 

96 yield cls._validate_tensor(spec, value) 

97 elif isinstance(spec, Mapping): 

98 if not isinstance(value, Mapping) or value.keys() != spec.keys(): 

99 raise ValueError("Batch fields do not match model_batch_spec.") 

100 for key, item in spec.items(): 

101 yield from cls._flatten(item, value[key]) 

102 elif isinstance(spec, (tuple, list)): 

103 if type(value) is not type(spec) or len(value) != len(spec): 

104 raise ValueError("Batch container does not match model_batch_spec.") 

105 for item, child in zip(spec, value): 

106 yield from cls._flatten(item, child) 

107 elif type(value) is not type(spec) or value != spec: 

108 raise ValueError("Non-tensor batch values must equal their static model_batch_spec values.") 

109 

110 @staticmethod 

111 def _validate_tensor(spec: TensorFieldSpec, value: Any) -> torch.Tensor: 

112 if not isinstance(value, torch.Tensor) or value.layout != torch.strided: 

113 raise ValueError("model_batch_spec requires a dense tensor at every TensorFieldSpec leaf.") 

114 if value.dtype != spec.dtype or value.ndim != len(spec.shape) or any( 

115 dim is not None and dim != actual for dim, actual in zip(spec.shape, value.shape) 

116 ): 

117 raise ValueError(f"Tensor shape/dtype {value.shape}/{value.dtype} does not match {spec}.") 

118 return value 

119 

120 def dimensions(self, tensors: list[torch.Tensor]) -> list[int]: 

121 """Extract only dimensions declared dynamic; fields and dtypes stay local. 

122 

123 Args: 

124 tensors: Validated source tensors in declaration order. 

125 """ 

126 return [actual for field, tensor in zip(self.fields, tensors) 

127 for dim, actual in zip(field.shape, tensor.shape) if dim is None] 

128 

129 def shapes(self, dimensions: list[int]) -> list[tuple[int, ...]]: 

130 """Resolve receive shapes from the compact dynamic-dimension header. 

131 

132 Args: 

133 dimensions: Dynamic sizes received in declaration order. 

134 """ 

135 if len(dimensions) != self.dynamic_dimensions or any(dim < 0 for dim in dimensions): 

136 raise ValueError("Invalid model batch dimension header.") 

137 pending = iter(dimensions) 

138 return [tuple(next(pending) if dim is None else dim for dim in field.shape) for field in self.fields] 

139 

140 def rebuild(self, tensors: list[torch.Tensor]) -> Any: 

141 """Restore the locally declared containers and static values. 

142 

143 Args: 

144 tensors: Broadcast tensor fields in declaration order. 

145 """ 

146 return self._rebuild(self.spec, iter(tensors)) 

147 

148 @classmethod 

149 def _rebuild(cls, spec: Any, tensors: Any) -> Any: 

150 if isinstance(spec, TensorFieldSpec): 

151 return next(tensors) 

152 if isinstance(spec, Mapping): 

153 return {key: cls._rebuild(item, tensors) for key, item in spec.items()} 

154 if isinstance(spec, (tuple, list)): 

155 return type(spec)(cls._rebuild(item, tensors) for item in spec) 

156 return spec