Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / batching / packing.py: 19%

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"""Select and combine samples for multimodal packing.""" 

16 

17from abc import ABC, abstractmethod 

18from collections.abc import Mapping, Sequence 

19from typing import Any 

20 

21import torch 

22 

23from hyper_parallel.data.constants import DEFAULT_FIELD_PACK_DIMS, PACKING_METADATA_FIELDS 

24from hyper_parallel.data.dataset_logging import get_dataset_logger 

25 

26logger = get_dataset_logger(__name__) 

27 

28 

29class PackingSelector(ABC): 

30 """Define the pre-packing seam over prepared sample metadata.""" 

31 

32 @abstractmethod 

33 def get_sample_cost(self, sample: Mapping[str, Any]) -> int: 

34 """Return the scheduling cost retained with one candidate sample.""" 

35 raise NotImplementedError 

36 

37 @abstractmethod 

38 def select_samples_to_pack( 

39 self, 

40 samples: Sequence[Mapping[str, Any]], 

41 token_budget: int, 

42 ) -> Sequence[int]: 

43 """Return indices of candidates that form the next packed sequence.""" 

44 raise NotImplementedError 

45 

46 

47class FirstFitPackingSelector(PackingSelector): 

48 """Select candidates in source order until the token budget is full.""" 

49 

50 def get_sample_cost(self, sample: Mapping[str, Any]) -> int: 

51 """Read standard packing metadata or derive cost from encoded input IDs.""" 

52 sample_cost = sample.get("packing_length") 

53 if sample_cost is None: 

54 sample_cost = self._get_input_length(sample) 

55 

56 resolved_cost = int(sample_cost) 

57 if resolved_cost <= 0: 

58 raise ValueError("Packing sample cost must be positive") 

59 

60 return resolved_cost 

61 

62 def select_samples_to_pack( 

63 self, 

64 samples: Sequence[Mapping[str, Any]], 

65 token_budget: int, 

66 ) -> Sequence[int]: 

67 """Select the next ordered first-fit group from the candidate buffer.""" 

68 if not samples: 

69 raise ValueError("Packing candidate buffer is empty") 

70 

71 selected_indices = [] 

72 selected_cost = 0 

73 for sample_index, sample in enumerate(samples): 

74 sample_cost = self.get_sample_cost(sample) 

75 sample_fits = selected_cost == 0 or selected_cost + sample_cost <= token_budget 

76 if sample_fits: 

77 selected_indices.append(sample_index) 

78 selected_cost += sample_cost 

79 

80 return selected_indices 

81 

82 @staticmethod 

83 def _get_input_length(sample: Mapping[str, Any]) -> int: 

84 """Derive sequence length from a fully encoded sample.""" 

85 if "input_ids" not in sample: 

86 raise ValueError("Packing samples must provide packing_length or input_ids") 

87 

88 input_ids = sample["input_ids"] 

89 input_shape = getattr(input_ids, "shape", None) 

90 if input_shape is not None: 

91 input_length = int(input_shape[-1]) 

92 return input_length 

93 

94 if not isinstance(input_ids, Sequence) or isinstance(input_ids, (str, bytes)): 

95 raise TypeError("Packing input_ids must be a tensor or sequence") 

96 

97 input_length = len(input_ids) 

98 if input_ids and isinstance(input_ids[0], Sequence): 

99 input_length = len(input_ids[0]) 

100 return input_length 

101 

102 

103class SamplePacker: 

104 """Combine selected encoded samples into one packed sample.""" 

105 

106 def __init__(self, field_pack_dims: Mapping[str, int] | None = None) -> None: 

107 """Configure the concatenation dimension of model fields. 

108 

109 Args: 

110 field_pack_dims: Optional additions or overrides to the default 

111 token-last and modality-first field rules. 

112 """ 

113 self.field_pack_dims = dict(DEFAULT_FIELD_PACK_DIMS) 

114 if field_pack_dims is not None: 

115 self.field_pack_dims.update(field_pack_dims) 

116 

117 def pack_selected_samples( 

118 self, 

119 samples: Sequence[Mapping[str, Any]], 

120 ) -> Mapping[str, Any]: 

121 """Concatenate one selected group and emit its sequence boundaries.""" 

122 if not samples: 

123 raise ValueError("samples must contain at least one encoded model sample") 

124 

125 packed_sample = {} 

126 sample_fields = set() 

127 for sample in samples: 

128 sample_fields.update(sample) 

129 

130 for field in sample_fields: 

131 if field in PACKING_METADATA_FIELDS: 

132 continue 

133 

134 if field.endswith("_token_starts"): 

135 packed_sample[field] = self._pack_token_starts(field, samples) 

136 continue 

137 

138 field_values = [] 

139 for sample in samples: 

140 if field in sample: 

141 field_values.append(sample[field]) 

142 

143 packed_value = self._pack_field(field, field_values) 

144 packed_sample[field] = packed_value 

145 

146 packed_sample["cu_seq_lens"] = self._build_sequence_boundaries(samples) 

147 return packed_sample 

148 

149 def _pack_field(self, field: str, values: Sequence[Any]) -> Any: 

150 """Pack one field using its configured or inferred leading dimension.""" 

151 if not values: 

152 raise ValueError(f"Cannot pack empty field {field!r}") 

153 

154 pack_dim = self.field_pack_dims.get(field, 0) 

155 if all(isinstance(value, torch.Tensor) for value in values): 

156 packed_value = torch.cat(values, dim=pack_dim) 

157 return packed_value 

158 

159 if field in self.field_pack_dims: 

160 tensor_values = [] 

161 for value in values: 

162 tensor_values.append(torch.as_tensor(value)) 

163 packed_value = torch.cat(tensor_values, dim=pack_dim) 

164 return packed_value 

165 

166 if all(isinstance(value, Sequence) and not isinstance(value, (str, bytes)) for value in values): 

167 packed_items = [] 

168 for value in values: 

169 packed_items.extend(value) 

170 return packed_items 

171 

172 if len(values) == 1: 

173 return values[0] 

174 

175 raise TypeError(f"Field {field!r} requires a model-specific SamplePacker rule") 

176 

177 @staticmethod 

178 def _pack_token_starts( 

179 field: str, 

180 samples: Sequence[Mapping[str, Any]], 

181 ) -> torch.Tensor: 

182 """Move sample-local modality starts into the packed token coordinate.""" 

183 packed_starts = [] 

184 sequence_offset = 0 

185 for sample in samples: 

186 if field in sample: 

187 token_starts = torch.as_tensor(sample[field]) 

188 packed_starts.append(token_starts + sequence_offset) 

189 

190 input_ids = sample.get("input_ids") 

191 if input_ids is None: 

192 raise ValueError(f"Packing {field!r} requires input_ids in every sample") 

193 sequence_offset += int(torch.as_tensor(input_ids).shape[-1]) 

194 

195 if not packed_starts: 

196 raise ValueError(f"Cannot pack empty field {field!r}") 

197 

198 packed_token_starts = torch.cat(packed_starts, dim=0) 

199 return packed_token_starts 

200 

201 @staticmethod 

202 def _build_sequence_boundaries( 

203 samples: Sequence[Mapping[str, Any]], 

204 ) -> torch.Tensor: 

205 """Build cumulative text boundaries for packed attention isolation.""" 

206 sequence_lengths = [] 

207 input_ids_reference = None 

208 for sample in samples: 

209 if "input_ids" not in sample: 

210 raise ValueError("Every packed sample must contain input_ids") 

211 

212 input_ids = sample["input_ids"] 

213 if input_ids_reference is None and isinstance(input_ids, torch.Tensor): 

214 input_ids_reference = input_ids 

215 input_shape = getattr(input_ids, "shape", None) 

216 if input_shape is not None: 

217 sequence_length = int(input_shape[-1]) 

218 else: 

219 sequence_length = len(input_ids) 

220 if sequence_length <= 0: 

221 raise ValueError("Packed samples must contain at least one token") 

222 sequence_lengths.append(sequence_length) 

223 

224 cumulative_lengths = [0] 

225 for sequence_length in sequence_lengths: 

226 cumulative_lengths.append(cumulative_lengths[-1] + sequence_length) 

227 

228 if input_ids_reference is None: 

229 sequence_boundaries = torch.tensor(cumulative_lengths, dtype=torch.int32) 

230 else: 

231 sequence_boundaries = input_ids_reference.new_tensor( 

232 cumulative_lengths, 

233 dtype=torch.int32, 

234 ) 

235 return sequence_boundaries 

236 

237 

238__all__ = ["FirstFitPackingSelector", "PackingSelector", "SamplePacker"]