Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / checkpoint / weight_conversion.py: 10%
252 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.
15# Adapted from:
16# https://github.com/huggingface/transformers/blob/v5.5.3/src/transformers/core_model_loading.py
17# https://github.com/huggingface/transformers/blob/v5.5.3/src/transformers/conversion_mapping.py
18# ============================================================================
19"""Transformers 4.57-and-later weight-conversion compatibility."""
21from __future__ import annotations
23__all__ = [
24 "ConversionOps",
25 "Transpose",
26 "WeightConverter",
27 "WeightRenaming",
28 "dot_natural_key",
29 "get_model_conversion_mapping",
30 "rename_source_key",
31 "revert_weight_conversion",
32]
34import re
35from abc import ABC, abstractmethod
36from collections import defaultdict
37from copy import deepcopy
38from typing import Any, Callable
40import torch
42try:
43 from transformers.conversion_mapping import (
44 get_model_conversion_mapping as _hf_get_model_conversion_mapping,
45 )
46 from transformers.core_model_loading import (
47 ConversionOps as _HFConversionOps,
48 Transpose as _HFTranspose,
49 WeightConverter as _HFWeightConverter,
50 WeightRenaming as _HFWeightRenaming,
51 )
52except ImportError:
53 _hf_get_model_conversion_mapping = None
54 _HFConversionOps = None
55 _HFTranspose = None
56 _HFWeightConverter = None
57 _HFWeightRenaming = None
59try:
60 from transformers.core_model_loading import (
61 revert_weight_conversion as _hf_revert_weight_conversion,
62 )
63except ImportError:
64 _hf_revert_weight_conversion = None
67def _has_scoped_weight_transforms() -> bool:
68 """Return whether the installed Transformers exposes scoped transforms."""
69 if _HFWeightConverter is None:
70 return False
71 return any(
72 "scope_prefix" in getattr(base, "__slots__", ())
73 for base in _HFWeightConverter.__mro__
74 )
77if _has_scoped_weight_transforms() and _hf_revert_weight_conversion is not None:
78 ConversionOps = _HFConversionOps
79 Transpose = _HFTranspose
80 WeightConverter = _HFWeightConverter
81 WeightRenaming = _HFWeightRenaming
82 from transformers.core_model_loading import dot_natural_key, rename_source_key
84 get_model_conversion_mapping = _hf_get_model_conversion_mapping
85 revert_weight_conversion = _hf_revert_weight_conversion
86else:
87 _WeightConverterBase = _HFWeightConverter or object
88 _WeightRenamingBase = _HFWeightRenaming or object
90 def _process_target_pattern(pattern: str) -> tuple[str, str | None]:
91 """Normalize one target pattern for reversible conversion."""
92 pattern = pattern.removeprefix("^").removesuffix("$")
93 pattern = re.sub(r"\(\?.+?\)?\)", "", pattern)
94 pattern = pattern.replace(r"\.", ".")
95 capturing_group = re.search(r"\(.+?\)", pattern)
96 if capturing_group is None:
97 return pattern, None
98 captured = capturing_group.group(0)
99 return pattern.replace(captured, r"\1", 1), captured
102 def _process_source_pattern(source_pattern: str, target_pattern: str) -> str:
103 """Preserve source anchors when building a reverse transform."""
104 if target_pattern.startswith("^") and not source_pattern.startswith("^"):
105 source_pattern = f"^{source_pattern}"
106 if target_pattern.endswith("$") and not source_pattern.endswith("$"):
107 source_pattern = f"{source_pattern}$"
108 return source_pattern
111 class ConversionOps(ABC):
112 """Base operation matching the Transformers 5.x conversion contract."""
114 @abstractmethod
115 def convert(
116 self,
117 input_dict: dict[str, Any],
118 source_patterns: list[str],
119 target_patterns: list[str],
120 **kwargs: Any,
121 ) -> dict[str, torch.Tensor]:
122 """Convert collected checkpoint tensors."""
124 @property
125 def reverse_op(self) -> "ConversionOps":
126 """Return the inverse operation."""
127 raise NotImplementedError
130 class Transpose(ConversionOps):
131 """Transpose the first two configured dimensions."""
133 def __init__(self, dim0: int = 0, dim1: int = 1, check_dims: bool = False) -> None:
134 """Configure the transpose operation.
136 Args:
137 dim0: First dimension to transpose.
138 dim1: Second dimension to transpose.
139 check_dims: Skip transposition when the tensor already matches the target shape.
140 """
141 self.dim0 = dim0
142 self.dim1 = dim1
143 self.check_dims = check_dims
145 @torch.no_grad()
146 def convert(
147 self,
148 input_dict: dict[str, torch.Tensor | list[torch.Tensor]],
149 source_patterns: list[str],
150 target_patterns: list[str],
151 **kwargs: Any,
152 ) -> dict[str, torch.Tensor]:
153 """Transpose the collected tensor into the selected target layout."""
154 if len(input_dict) != 1:
155 raise ValueError("Transpose requires exactly one collected source")
156 if len(target_patterns) > 1:
157 if len(source_patterns) != 1:
158 raise ValueError("Transpose cannot select an intermediate pattern")
159 target_pattern = source_patterns[0]
160 else:
161 target_pattern = target_patterns[0]
162 value = next(iter(input_dict.values()))
163 tensor = value[0] if isinstance(value, list) else value
164 if self.check_dims:
165 model = kwargs["model"]
166 expected = model.get_parameter(kwargs["full_layer_name"]).shape
167 if tensor.shape == expected:
168 return {target_pattern: tensor}
169 return {
170 target_pattern: torch.transpose(
171 tensor, dim0=self.dim0, dim1=self.dim1
172 ).contiguous()
173 }
175 @property
176 def reverse_op(self) -> "ConversionOps":
177 """Return a transpose with the same dimensions to undo this operation."""
178 return Transpose(self.dim0, self.dim1, self.check_dims)
181 class _WeightTransform:
182 """Minimal Transformers 5.x weight-transform contract for 4.57."""
184 def __init__(
185 self,
186 source_patterns: str | list[str],
187 target_patterns: str | list[str],
188 ) -> None:
189 """Normalize the source and target patterns for reversible conversion.
191 Args:
192 source_patterns: Checkpoint key patterns collected by this transform.
193 target_patterns: Model key patterns produced by this transform.
194 """
195 self.source_patterns = (
196 [source_patterns] if isinstance(source_patterns, str) else list(source_patterns)
197 )
198 self.target_patterns = (
199 [target_patterns] if isinstance(target_patterns, str) else list(target_patterns)
200 )
201 self._original_source_patterns = self.source_patterns.copy()
202 self._original_target_patterns = self.target_patterns.copy()
203 captured_groups = []
204 for index, pattern in enumerate(self.target_patterns):
205 self.target_patterns[index], captured = _process_target_pattern(pattern)
206 if captured is not None:
207 captured_groups.append(captured)
208 unique_captured_groups = set(captured_groups)
209 if len(unique_captured_groups) > 1:
210 raise ValueError(
211 "target_patterns must use at most one unique capturing group"
212 )
213 captured_group = (
214 unique_captured_groups.pop() if unique_captured_groups else None
215 )
216 for index, pattern in enumerate(self.source_patterns):
217 if r"\1" in pattern:
218 if captured_group is None:
219 raise ValueError(
220 "source pattern contains a backreference without a "
221 "target capturing group"
222 )
223 pattern = pattern.replace(r"\1", captured_group, 1)
224 if len(self.source_patterns) == len(self.target_patterns):
225 pattern = _process_source_pattern(
226 pattern, self._original_target_patterns[index]
227 )
228 self.source_patterns[index] = pattern
229 self.collected_tensors: dict[
230 str, list[torch.Tensor | Callable[[], torch.Tensor]]
231 ] = defaultdict(list)
232 self.scope_prefix: str | None = None
233 self.base_model_prefix: str | None = None
234 self._was_used = False
235 branches = []
236 for index, pattern in enumerate(self.source_patterns):
237 pattern = pattern.replace(".*.", r"\..*\.")
238 branches.append(f"(?P<g{index}>{pattern})")
239 self._compiled_sources = re.compile("|".join(branches))
241 def _key_to_match(self, source_key: str) -> tuple[str | None, str] | None:
242 if self.scope_prefix is None:
243 return None, source_key
244 scope = f"{self.scope_prefix}." if self.scope_prefix else ""
245 base = f"{self.base_model_prefix}." if self.base_model_prefix else ""
246 for prefix in (base + scope, scope):
247 if source_key.startswith(prefix):
248 return prefix, source_key.removeprefix(prefix)
249 return None
251 def rename_source_key(self, source_key: str) -> tuple[str, str | None]:
252 """Return the renamed key and the source pattern that matched it."""
253 scoped_key = self._key_to_match(source_key)
254 if scoped_key is None:
255 return source_key, None
256 prefix, key_to_match = scoped_key
257 match = self._compiled_sources.search(key_to_match)
258 if match is None:
259 return source_key, None
260 group_name = next(
261 name for name, value in match.groupdict().items() if value is not None
262 )
263 pattern = self.source_patterns[int(group_name[1:])]
264 replacement = self.target_patterns[0]
265 if r"\1" in replacement:
266 group_index = self._compiled_sources.groupindex[group_name] + 1
267 replacement = replacement.replace(
268 r"\1", match.group(group_index), 1
269 )
270 self._was_used = True
271 renamed = key_to_match.replace(match.group(0), replacement, 1)
272 return f"{prefix or ''}{renamed}", pattern
274 def add_tensor(
275 self,
276 target_key: str,
277 source_key: str,
278 source_pattern: str,
279 loader: Callable[[], torch.Tensor],
280 ) -> None:
281 """Collect a lazy checkpoint tensor under its matched source pattern.
283 Args:
284 target_key: Target checkpoint key accepted for API compatibility.
285 source_key: Source checkpoint key accepted for API compatibility.
286 source_pattern: Pattern used to group tensors for conversion.
287 loader: Callable that materializes the source tensor.
288 """
289 del target_key, source_key
290 self.collected_tensors[source_pattern].append(loader)
292 def _materialize_tensors(self) -> dict[str, list[torch.Tensor]]:
293 materialized = {}
294 for pattern, values in list(self.collected_tensors.items()):
295 materialized[pattern] = [
296 value() if callable(value) else value for value in values
297 ]
298 self.collected_tensors.clear()
299 return materialized
301 def was_used(self) -> bool:
302 """Return whether this transform matched a checkpoint key."""
303 return self._was_used
305 def reverse_transform(self) -> "_WeightTransform":
306 """Build the transform used to restore original checkpoint names."""
307 kwargs = {}
308 if hasattr(self, "operations"):
309 kwargs["operations"] = [
310 operation.reverse_op for operation in self.operations[::-1]
311 ]
312 transform = self.__class__(
313 source_patterns=self._original_target_patterns,
314 target_patterns=self._original_source_patterns,
315 **kwargs,
316 )
317 transform.scope_prefix = self.scope_prefix
318 transform.base_model_prefix = self.base_model_prefix
319 return transform
322 class WeightRenaming(_WeightTransform, _WeightRenamingBase):
323 """Rename one checkpoint pattern without changing its tensor."""
325 def convert(
326 self,
327 layer_name: str,
328 model: Any = None,
329 config: Any = None,
330 hf_quantizer: Any = None,
331 loading_info: Any = None,
332 ) -> dict[str, torch.Tensor]:
333 """Materialize and rename the collected checkpoint tensor."""
334 del model, config, hf_quantizer, loading_info
335 collected = self._materialize_tensors()
336 tensor = collected[self.source_patterns[0]]
337 target = self.target_patterns[0]
338 if target not in layer_name:
339 target = layer_name
340 return {target: tensor}
343 class WeightConverter(_WeightTransform, _WeightConverterBase):
344 """Apply conversion operations to one checkpoint tensor group."""
346 def __init__(
347 self,
348 source_patterns: str | list[str],
349 target_patterns: str | list[str],
350 operations: list[ConversionOps],
351 ) -> None:
352 """Configure the operations applied to a collected checkpoint group.
354 Args:
355 source_patterns: Checkpoint key patterns collected by this converter.
356 target_patterns: Model key patterns produced by this converter.
357 operations: Nonempty sequence of conversion operations, applied in order.
358 """
359 super().__init__(source_patterns, target_patterns)
360 if not operations:
361 raise ValueError("WeightConverter requires at least one operation")
362 self.operations = operations
364 def convert(
365 self,
366 layer_name: str,
367 model: Any = None,
368 config: Any = None,
369 hf_quantizer: Any = None,
370 loading_info: Any = None,
371 ) -> dict[str, torch.Tensor]:
372 """Materialize the collected tensors and apply each operation."""
373 collected: dict[str, Any] = self._materialize_tensors()
374 for operation in self.operations:
375 collected = operation.convert(
376 collected,
377 source_patterns=self.source_patterns,
378 target_patterns=self.target_patterns,
379 full_layer_name=layer_name,
380 model=model,
381 config=config,
382 missing_keys=(
383 loading_info.missing_keys if loading_info is not None else None
384 ),
385 )
386 del hf_quantizer
387 full_name = layer_name.replace(".*.", ".0.")
388 try:
389 prefix, _, suffix = next(
390 full_name.partition(pattern)
391 for pattern in collected
392 if pattern in full_name
393 )
394 except StopIteration:
395 return collected
396 return {f"{prefix}{name}{suffix}": value for name, value in collected.items()}
399 def rename_source_key(
400 source_key: str,
401 weight_renamings: list[WeightRenaming],
402 weight_converters: list[WeightConverter],
403 base_model_prefix: str | None = None,
404 meta_state_dict: dict[str, Any] | None = None,
405 ) -> tuple[str, str | None]:
406 """Apply renamings, one converter, and base-prefix normalization."""
407 renamed_key = source_key
408 for renaming in weight_renamings:
409 renamed_key, _ = renaming.rename_source_key(renamed_key)
410 source_pattern = None
411 for converter in weight_converters:
412 renamed_key, source_pattern = converter.rename_source_key(renamed_key)
413 if source_pattern is not None:
414 break
415 if base_model_prefix is not None and meta_state_dict is not None:
416 prefix = f"{base_model_prefix}."
417 without_prefix = renamed_key.removeprefix(prefix)
418 if renamed_key.startswith(prefix) and without_prefix in meta_state_dict:
419 renamed_key = without_prefix
420 elif f"{prefix}{renamed_key}" in meta_state_dict:
421 renamed_key = f"{prefix}{renamed_key}"
422 return renamed_key, source_pattern
425 def dot_natural_key(value: str) -> list[tuple[int, int | str]]:
426 """Sort dotted tensor names with numeric path components numerically."""
427 return [
428 (0, int(part)) if part.isdigit() else (1, part)
429 for part in value.split(".")
430 ]
433 def get_model_conversion_mapping(
434 model: Any,
435 key_mapping: dict[str, str] | None = None,
436 hf_quantizer: Any = None,
437 add_legacy: bool = True,
438 ) -> list[WeightRenaming | WeightConverter]:
439 """Normalize available Transformers mappings to scoped transforms."""
440 if _hf_get_model_conversion_mapping is None:
441 return [
442 WeightRenaming(source_patterns=source, target_patterns=target)
443 for source, target in (key_mapping or {}).items()
444 ]
445 transforms = _hf_get_model_conversion_mapping(
446 model,
447 key_mapping=key_mapping,
448 hf_quantizer=hf_quantizer,
449 add_legacy=add_legacy,
450 )
451 normalized = []
452 for transform in transforms:
453 if isinstance(transform, _HFWeightConverter):
454 normalized.append(
455 WeightConverter(
456 source_patterns=transform.source_patterns,
457 target_patterns=transform.target_patterns,
458 operations=deepcopy(transform.operations),
459 )
460 )
461 elif isinstance(transform, _HFWeightRenaming):
462 normalized.append(
463 WeightRenaming(
464 source_patterns=transform.source_patterns,
465 target_patterns=transform.target_patterns,
466 )
467 )
468 else:
469 normalized.append(transform)
470 return normalized
473 def _collect_reverse_conversions(
474 state_dict: dict[str, torch.Tensor],
475 renamings: list[WeightRenaming],
476 converters: list[WeightConverter],
477 ) -> dict[str, _WeightTransform]:
478 """Collect reverse transforms and their source tensors."""
479 conversion_mapping = {}
480 for original_key, tensor in sorted(
481 state_dict.items(), key=lambda item: dot_natural_key(item[0])
482 ):
483 renamed_key = original_key
484 for renaming in renamings:
485 renamed_key, _ = renaming.rename_source_key(renamed_key)
487 source_pattern = None
488 matched_converter = None
489 for converter in converters:
490 renamed_key, source_pattern = converter.rename_source_key(renamed_key)
491 if source_pattern is not None:
492 matched_converter = converter
493 break
495 if matched_converter is not None:
496 transform = conversion_mapping.setdefault(
497 renamed_key, deepcopy(matched_converter)
498 )
499 else:
500 transform = conversion_mapping.setdefault(
501 renamed_key, WeightRenaming(original_key, renamed_key)
502 )
503 source_pattern = original_key
504 transform.add_tensor(renamed_key, original_key, source_pattern, tensor)
505 return conversion_mapping
508 def _materialize_reverse_conversions(
509 model: Any,
510 conversion_mapping: dict[str, _WeightTransform],
511 ) -> dict[str, torch.Tensor]:
512 """Apply collected reverse transforms to produce checkpoint tensors."""
513 converted_state_dict = {}
514 for first_param_name, reverse_conversion in conversion_mapping.items():
515 realized = reverse_conversion.convert(
516 first_param_name, model=model, config=model.config
517 )
518 for target_name, parameter in realized.items():
519 converted_state_dict[target_name] = (
520 parameter[0] if isinstance(parameter, list) else parameter
521 )
522 return converted_state_dict
525 def revert_weight_conversion(
526 model: Any,
527 state_dict: dict[str, torch.Tensor],
528 ) -> dict[str, torch.Tensor]:
529 """Restore checkpoint names and layouts used before weight conversion."""
530 weight_conversions = getattr(model, "_weight_conversions", None)
531 if weight_conversions is None:
532 weight_conversions = get_model_conversion_mapping(model, add_legacy=False)
533 if not weight_conversions:
534 return state_dict
536 reverse_conversions = [
537 conversion.reverse_transform() for conversion in weight_conversions
538 ]
539 renamings = [
540 conversion
541 for conversion in reverse_conversions
542 if isinstance(conversion, WeightRenaming)
543 ]
544 converters = [
545 conversion
546 for conversion in reverse_conversions
547 if isinstance(conversion, WeightConverter)
548 ]
549 conversion_mapping = _collect_reverse_conversions(
550 state_dict, renamings, converters
551 )
552 return _materialize_reverse_conversions(model, conversion_mapping)