Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / checkpoint / conversion_ops.py: 86%
169 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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"""Reusable checkpoint conversion operations."""
17from typing import Any
19import torch # pylint: disable=forbidden-backend-import
20from hyper_parallel.components.checkpoint.weight_conversion import ConversionOps
23def _single_tensor(value: torch.Tensor | list[torch.Tensor]) -> torch.Tensor:
24 """Return the single tensor collected for one checkpoint pattern."""
25 if isinstance(value, list):
26 if len(value) != 1:
27 raise ValueError("expected exactly one tensor per checkpoint pattern")
28 return value[0]
29 return value
32class ConcatenateWithSections(ConversionOps):
33 """Concatenate tensors and retain their section sizes for reverse conversion."""
35 def __init__(self, sections: tuple[int, ...], dim: int = 0) -> None:
36 """Configure the tensor sizes along the concatenation dimension."""
37 if len(sections) < 2 or any(not isinstance(size, int) or size <= 0 for size in sections):
38 raise ValueError("sections must contain at least two positive integers")
39 self.sections = tuple(sections)
40 self.dim = dim
42 @torch.no_grad()
43 def convert(
44 self,
45 input_dict: dict[str, torch.Tensor | list[torch.Tensor]],
46 source_patterns: list[str],
47 target_patterns: list[str],
48 **kwargs: Any,
49 ) -> dict[str, torch.Tensor]:
50 """Concatenate one tensor from each source pattern."""
51 del kwargs
52 if len(source_patterns) != len(self.sections) or len(target_patterns) != 1:
53 raise ValueError(
54 "ConcatenateWithSections requires one source per section and one target"
55 )
56 tensors = [_single_tensor(input_dict[pattern]) for pattern in source_patterns]
57 actual_sections = tuple(tensor.shape[self.dim] for tensor in tensors)
58 if actual_sections != self.sections:
59 raise ValueError(
60 f"source sizes {actual_sections} do not match configured sections {self.sections}"
61 )
62 return {target_patterns[0]: torch.cat(tensors, dim=self.dim).contiguous()}
64 @property
65 def reverse_op(self) -> ConversionOps:
66 """Return the section-based split used for reverse conversion."""
67 return Split(self.sections, dim=self.dim)
70class Split(ConversionOps):
71 """Split a tensor into explicitly sized sections."""
73 def __init__(self, sections: tuple[int, ...], dim: int = 0) -> None:
74 """Configure the output sizes along the split dimension."""
75 if len(sections) < 2 or any(not isinstance(size, int) or size <= 0 for size in sections):
76 raise ValueError("sections must contain at least two positive integers")
77 self.sections = tuple(sections)
78 self.dim = dim
80 @torch.no_grad()
81 def convert(
82 self,
83 input_dict: dict[str, torch.Tensor | list[torch.Tensor]],
84 source_patterns: list[str],
85 target_patterns: list[str],
86 **kwargs: Any,
87 ) -> dict[str, torch.Tensor]:
88 """Split one source tensor and assign each section to one target pattern."""
89 del kwargs
90 if len(source_patterns) != 1 or len(target_patterns) != len(self.sections):
91 raise ValueError("Split requires one source and one target per section")
92 tensor = _single_tensor(input_dict[source_patterns[0]])
93 tensors = torch.split(tensor, self.sections, dim=self.dim)
94 return {
95 pattern: value.contiguous()
96 for pattern, value in zip(target_patterns, tensors)
97 }
99 @property
100 def reverse_op(self) -> ConversionOps:
101 """Return the section-aware concatenation used for reverse conversion."""
102 return ConcatenateWithSections(self.sections, dim=self.dim)
105class AddScalar(ConversionOps):
106 """Add a scalar to one checkpoint tensor."""
108 def __init__(self, value: float) -> None:
109 """Configure the scalar added during conversion."""
110 self.value = value
112 @torch.no_grad()
113 def convert(
114 self,
115 input_dict: dict[str, torch.Tensor | list[torch.Tensor]],
116 source_patterns: list[str],
117 target_patterns: list[str],
118 **kwargs: Any,
119 ) -> dict[str, torch.Tensor]:
120 """Add the configured scalar and assign the converted tensor."""
121 del kwargs
122 if len(source_patterns) != 1 or len(target_patterns) != 1:
123 raise ValueError("AddScalar requires exactly one source and one target")
124 tensor = _single_tensor(input_dict[source_patterns[0]])
125 return {target_patterns[0]: tensor + self.value}
127 @property
128 def reverse_op(self) -> ConversionOps:
129 """Return the inverse scalar addition used during export."""
130 return AddScalar(-self.value)
133class InterleaveQKV(ConversionOps):
134 """Convert Q, K, and V rows to the GQA-interleaved layout."""
136 def __init__(
137 self,
138 num_key_value_heads: int,
139 num_key_value_groups: int,
140 query_head_dim: int,
141 value_head_dim: int,
142 source_is_fused: bool,
143 ) -> None:
144 """Configure the source projection layout and GQA dimensions."""
145 self.num_key_value_heads = num_key_value_heads
146 self.num_key_value_groups = num_key_value_groups
147 self.query_head_dim = query_head_dim
148 self.value_head_dim = value_head_dim
149 self.source_is_fused = source_is_fused
151 @property
152 def _source_sections(self) -> tuple[int, int, int]:
153 return (
154 self.num_key_value_heads
155 * self.num_key_value_groups
156 * self.query_head_dim,
157 self.num_key_value_heads * self.query_head_dim,
158 self.num_key_value_heads * self.value_head_dim,
159 )
161 @torch.no_grad()
162 def convert(
163 self,
164 input_dict: dict[str, torch.Tensor | list[torch.Tensor]],
165 source_patterns: list[str],
166 target_patterns: list[str],
167 **kwargs: Any,
168 ) -> dict[str, torch.Tensor]:
169 """Convert fused or separate Q/K/V rows to the grouped layout."""
170 del kwargs
171 if len(target_patterns) != 1:
172 raise ValueError("grouped GQA projection requires exactly one target")
173 if self.source_is_fused:
174 if len(source_patterns) != 1:
175 raise ValueError("fused GQA projection requires exactly one source")
176 fused = _single_tensor(input_dict[source_patterns[0]])
177 query, key, value = torch.split(fused, self._source_sections, dim=0)
178 else:
179 if len(source_patterns) != 3:
180 raise ValueError("separate GQA projection requires Q, K, and V sources")
181 query, key, value = (
182 _single_tensor(input_dict[pattern]) for pattern in source_patterns
183 )
185 trailing_shape = query.shape[1:]
186 query = query.reshape(
187 self.num_key_value_heads,
188 self.num_key_value_groups * self.query_head_dim,
189 *trailing_shape,
190 )
191 key = key.reshape(
192 self.num_key_value_heads,
193 self.query_head_dim,
194 *trailing_shape,
195 )
196 value = value.reshape(
197 self.num_key_value_heads,
198 self.value_head_dim,
199 *trailing_shape,
200 )
201 grouped = torch.cat((query, key, value), dim=1).flatten(0, 1)
202 return {target_patterns[0]: grouped.contiguous()}
204 @property
205 def reverse_op(self) -> ConversionOps:
206 """Return the conversion that restores the source Q/K/V layout."""
207 return DeinterleaveQKV(
208 self.num_key_value_heads,
209 self.num_key_value_groups,
210 self.query_head_dim,
211 self.value_head_dim,
212 self.source_is_fused,
213 )
216class DeinterleaveQKV(ConversionOps):
217 """Restore concatenated or separate Q, K, and V rows."""
219 def __init__(
220 self,
221 num_key_value_heads: int,
222 num_key_value_groups: int,
223 query_head_dim: int,
224 value_head_dim: int,
225 target_is_fused: bool,
226 ) -> None:
227 """Configure the grouped layout and restored projection form."""
228 self.num_key_value_heads = num_key_value_heads
229 self.num_key_value_groups = num_key_value_groups
230 self.query_head_dim = query_head_dim
231 self.value_head_dim = value_head_dim
232 self.target_is_fused = target_is_fused
234 @property
235 def _group_sections(self) -> tuple[int, int, int]:
236 return (
237 self.num_key_value_groups * self.query_head_dim,
238 self.query_head_dim,
239 self.value_head_dim,
240 )
242 @torch.no_grad()
243 def convert(
244 self,
245 input_dict: dict[str, torch.Tensor | list[torch.Tensor]],
246 source_patterns: list[str],
247 target_patterns: list[str],
248 **kwargs: Any,
249 ) -> dict[str, torch.Tensor]:
250 """Restore one grouped projection as fused or separate Q/K/V rows."""
251 del kwargs
252 if len(source_patterns) != 1:
253 raise ValueError("ungrouping GQA projection requires exactly one source")
254 grouped = _single_tensor(input_dict[source_patterns[0]])
255 grouped = grouped.reshape(
256 self.num_key_value_heads,
257 sum(self._group_sections),
258 *grouped.shape[1:],
259 )
260 query, key, value = torch.split(grouped, self._group_sections, dim=1)
261 query = query.flatten(0, 1).contiguous()
262 key = key.flatten(0, 1).contiguous()
263 value = value.flatten(0, 1).contiguous()
265 if self.target_is_fused:
266 if len(target_patterns) != 1:
267 raise ValueError("fused GQA projection requires exactly one target")
268 return {target_patterns[0]: torch.cat((query, key, value), dim=0)}
269 if len(target_patterns) != 3:
270 raise ValueError("separate GQA projection requires Q, K, and V targets")
271 return dict(zip(target_patterns, (query, key, value)))
273 @property
274 def reverse_op(self) -> ConversionOps:
275 """Return the conversion that groups Q/K/V rows by KV head."""
276 return InterleaveQKV(
277 self.num_key_value_heads,
278 self.num_key_value_groups,
279 self.query_head_dim,
280 self.value_head_dim,
281 self.target_is_fused,
282 )
285class InterleaveGateQKV(ConversionOps):
286 """Group per-query-head Q/gate plus K/V rows into one projection."""
288 def __init__(
289 self,
290 num_key_value_heads: int,
291 num_key_value_groups: int,
292 query_head_dim: int,
293 value_head_dim: int,
294 ) -> None:
295 """Configure the gated query layout and GQA dimensions."""
296 self.num_key_value_heads = num_key_value_heads
297 self.num_key_value_groups = num_key_value_groups
298 self.query_head_dim = query_head_dim
299 self.value_head_dim = value_head_dim
301 @torch.no_grad()
302 def convert(
303 self,
304 input_dict: dict[str, torch.Tensor | list[torch.Tensor]],
305 source_patterns: list[str],
306 target_patterns: list[str],
307 **kwargs: Any,
308 ) -> dict[str, torch.Tensor]:
309 """Convert interleaved Q/gate plus K/V into grouped Q/K/V/gate rows."""
310 del kwargs
311 if len(source_patterns) != 3 or len(target_patterns) != 1:
312 raise ValueError(
313 "gated GQA conversion requires Q/gate, K, V sources and one target"
314 )
315 query_gate, key, value = (
316 _single_tensor(input_dict[pattern]) for pattern in source_patterns
317 )
318 trailing_shape = query_gate.shape[1:]
319 query_gate = query_gate.reshape(
320 self.num_key_value_heads,
321 self.num_key_value_groups,
322 2,
323 self.query_head_dim,
324 *trailing_shape,
325 )
326 query = query_gate[:, :, 0].reshape(
327 self.num_key_value_heads,
328 self.num_key_value_groups * self.query_head_dim,
329 *trailing_shape,
330 )
331 gate = query_gate[:, :, 1].reshape(
332 self.num_key_value_heads,
333 self.num_key_value_groups * self.query_head_dim,
334 *trailing_shape,
335 )
336 key = key.reshape(
337 self.num_key_value_heads,
338 self.query_head_dim,
339 *trailing_shape,
340 )
341 value = value.reshape(
342 self.num_key_value_heads,
343 self.value_head_dim,
344 *trailing_shape,
345 )
346 grouped = torch.cat((query, key, value, gate), dim=1).flatten(0, 1)
347 return {target_patterns[0]: grouped.contiguous()}
349 @property
350 def reverse_op(self) -> ConversionOps:
351 """Return the conversion that restores per-head Q/gate rows and K/V."""
352 return DeinterleaveGateQKV(
353 self.num_key_value_heads,
354 self.num_key_value_groups,
355 self.query_head_dim,
356 self.value_head_dim,
357 )
360class DeinterleaveGateQKV(ConversionOps):
361 """Restore per-query-head Q/gate rows and separate K/V rows."""
363 def __init__(
364 self,
365 num_key_value_heads: int,
366 num_key_value_groups: int,
367 query_head_dim: int,
368 value_head_dim: int,
369 ) -> None:
370 """Configure the grouped gated-GQA layout."""
371 self.num_key_value_heads = num_key_value_heads
372 self.num_key_value_groups = num_key_value_groups
373 self.query_head_dim = query_head_dim
374 self.value_head_dim = value_head_dim
376 @property
377 def _group_sections(self) -> tuple[int, int, int, int]:
378 return (
379 self.num_key_value_groups * self.query_head_dim,
380 self.query_head_dim,
381 self.value_head_dim,
382 self.num_key_value_groups * self.query_head_dim,
383 )
385 @torch.no_grad()
386 def convert(
387 self,
388 input_dict: dict[str, torch.Tensor | list[torch.Tensor]],
389 source_patterns: list[str],
390 target_patterns: list[str],
391 **kwargs: Any,
392 ) -> dict[str, torch.Tensor]:
393 """Restore Q/gate, K, and V rows from one grouped projection."""
394 del kwargs
395 if len(source_patterns) != 1 or len(target_patterns) != 3:
396 raise ValueError(
397 "gated GQA reverse conversion requires one source and Q/gate, K, V targets"
398 )
399 grouped = _single_tensor(input_dict[source_patterns[0]])
400 trailing_shape = grouped.shape[1:]
401 grouped = grouped.reshape(
402 self.num_key_value_heads,
403 sum(self._group_sections),
404 *trailing_shape,
405 )
406 query, key, value, gate = torch.split(
407 grouped, self._group_sections, dim=1
408 )
409 query = query.reshape(
410 self.num_key_value_heads,
411 self.num_key_value_groups,
412 self.query_head_dim,
413 *trailing_shape,
414 )
415 gate = gate.reshape(
416 self.num_key_value_heads,
417 self.num_key_value_groups,
418 self.query_head_dim,
419 *trailing_shape,
420 )
421 query_gate = torch.stack((query, gate), dim=2).flatten(0, 3)
422 return {
423 target_patterns[0]: query_gate.contiguous(),
424 target_patterns[1]: key.flatten(0, 1).contiguous(),
425 target_patterns[2]: value.flatten(0, 1).contiguous(),
426 }
428 @property
429 def reverse_op(self) -> ConversionOps:
430 """Return the conversion that groups Q/K/V and extracts gate rows."""
431 return InterleaveGateQKV(
432 self.num_key_value_heads,
433 self.num_key_value_groups,
434 self.query_head_dim,
435 self.value_head_dim,
436 )