Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / offline_transform.py: 34%
190 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd. All rights reserved.
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"""Offline checkpoint conversion: Hugging Face safetensors layout and HyperParallel DCP.
17Covers **full checkpoint** (unsharded weights) and **shard checkpoint** (e.g. HF multi-file safetensors),
18``state_dict`` validation, HF-style save/load, and DCP read/write (see ``.agent/rules/offline_transform.md``).
20Public symbols are listed in ``__all__``; sharding helpers and split types are internal.
21"""
22from __future__ import annotations
24import json
25import logging
26import os
27from dataclasses import dataclass
28from typing import Any, Literal
30import torch
32from hyper_parallel.core.distributed_checkpoint.filesystem_storage import (
33 FileSystemReader,
34 FileSystemWriter,
35)
36from hyper_parallel.core.distributed_checkpoint.standard_planner import (
37 StandardSavePlanner,
38 _DcpMergeLoadPlanner,
39)
40from hyper_parallel.core.distributed_checkpoint.checkpoint_io import (
41 load_checkpoint_file,
42 save_checkpoint_file,
43)
44from hyper_parallel.core.distributed_checkpoint.utils import get_tensor_storage_size
46logger = logging.getLogger(__name__)
48__all__ = [
49 "save_state_dict_as_huggingface_format",
50 "parse_checkpoint_from_huggingface",
51 "full_state_dict_to_dcp_format",
52 "dcp_to_full_state_dict",
53 "convert_full_checkpoint_to_dcp",
54]
57def _validate_state_dict(state_dict: dict[str, Any]) -> None:
58 """Ensure values are allowed before DCP / HF export.
60 Applies to a full-weights ``state_dict`` (from a full checkpoint file or merged shard checkpoint content).
62 Allowed values are ``torch.Tensor`` (including ``nn.Parameter``), ``bytes`` (e.g. optimizer blobs), or
63 ``str`` placeholders (e.g. bitsandbytes serialization per HF ecosystem).
65 Args:
66 state_dict: Flat or nested mapping rejected here if values are not allowed types.
68 Raises:
69 TypeError: If a value has an unsupported type.
70 """
71 if not state_dict:
72 return
73 for key, value in state_dict.items():
74 if isinstance(value, (torch.Tensor, bytes, str)):
75 continue
76 logger.warning(
77 "Unsupported offline checkpoint value type for key %r: %s "
78 "(expected torch.Tensor, bytes, or str per project rules).",
79 key,
80 type(value).__name__,
81 )
82 raise TypeError(
83 f"Offline checkpoint expects torch.Tensor, bytes, or str per key; "
84 f"got {type(value).__name__} for key {key!r}."
85 )
88_SIZE_UNITS = {
89 "TB": 10**12,
90 "GB": 10**9,
91 "MB": 10**6,
92 "KB": 10**3,
93}
95_SAFE_WEIGHTS_NAME = "model.safetensors"
96_SAFE_WEIGHTS_INDEX_NAME = "model.safetensors.index.json"
97# Hugging Face-style shard filenames: ``model.safetensors`` or ``model-00001-of-00002.safetensors``.
98_HF_SAFE_WEIGHTS_FILENAME_PATTERN = _SAFE_WEIGHTS_NAME.replace(".safetensors", "{suffix}.safetensors")
101@dataclass
102class _StateDictSplit:
103 """Result of splitting a state dict into named shard files and an optional weight map."""
105 metadata: dict[str, Any]
106 filename_to_tensors: dict[str, list[str]]
107 tensor_to_filename: dict[str, str]
108 is_sharded: bool = False
110 def __post_init__(self) -> None:
111 self.is_sharded = len(self.filename_to_tensors) > 1
114def _parse_size_to_int(size_as_str: str) -> int:
115 """
116 Parse a size expressed as a string with digits and unit (like `"5MB"`) to an integer (in bytes).
118 Supported units are "TB", "GB", "MB", "KB".
120 Args:
121 size_as_str (`str`): The size to convert. Will be directly returned if an `int`.
123 Example:
125 ```py
126 >>> _parse_size_to_int("5MB")
127 5000000
128 ```
129 """
130 size_as_str = size_as_str.strip()
132 # Parse unit
133 unit = size_as_str[-2:].upper()
134 if unit not in _SIZE_UNITS:
135 raise ValueError(f"Unit '{unit}' not supported. Supported units are TB, GB, MB, KB. Got '{size_as_str}'.")
136 multiplier = _SIZE_UNITS[unit]
138 # Parse value
139 try:
140 value = float(size_as_str[:-2].strip())
141 except ValueError as e:
142 raise ValueError(f"Could not parse the size value from '{size_as_str}': {e}") from e
144 return int(value * multiplier)
147def _build_shard_list_for_safetensors(
148 state_dict: dict[str, Any],
149 max_cap: int,
150 get_storage_size: Any,
151) -> list[dict[str, Any]]:
152 """Assign tensors to shards in key order (greedy by ``max_cap`` bytes)."""
153 shard_list: list[dict[str, Any]] = []
154 current_shard: dict[str, Any] = {}
155 current_shard_size = 0
157 for key, tensor in state_dict.items():
158 tensor_size = get_storage_size(tensor)
160 if tensor_size > max_cap:
161 shard_list.append({key: tensor})
162 continue
164 if current_shard_size + tensor_size > max_cap:
165 shard_list.append(current_shard)
166 current_shard = {}
167 current_shard_size = 0
169 current_shard[key] = tensor
170 current_shard_size += tensor_size
172 if len(current_shard) > 0:
173 shard_list.append(current_shard)
174 return shard_list
177def _total_bytes_unique_keys_in_shards(
178 shard_list: list[dict[str, Any]],
179 get_storage_size: Any,
180) -> int:
181 """Sum storage bytes once per parameter key (excludes str placeholders)."""
182 total_size = 0
183 seen_keys: set[str] = set()
184 for shard in shard_list:
185 for k, tensor in shard.items():
186 if k in seen_keys:
187 continue
188 seen_keys.add(k)
189 if isinstance(tensor, str):
190 continue
191 total_size += get_storage_size(tensor)
192 return total_size
195def _split_state_dict_into_shards(
196 state_dict: dict[str, Any],
197 max_shard_size: int | str,
198) -> _StateDictSplit:
199 """Shard ``state_dict`` for safetensors by serialized tensor byte size.
201 Shards are built by iterating ``state_dict`` in key order (no bin-packing optimization). If a single tensor exceeds
202 ``max_shard_size``, it occupies its own shard (possibly larger than the cap).
204 Shard filenames follow Hugging Face conventions (``model.safetensors`` or
205 ``model-00001-of-00002.safetensors``).
206 """
207 get_storage_size = get_tensor_storage_size
209 filename_pattern = _HF_SAFE_WEIGHTS_FILENAME_PATTERN
211 max_cap: int | str = max_shard_size
212 if isinstance(max_cap, str):
213 max_cap = _parse_size_to_int(max_cap)
215 shard_list = _build_shard_list_for_safetensors(state_dict, max_cap, get_storage_size)
216 nb_shards = len(shard_list)
218 total_size = _total_bytes_unique_keys_in_shards(shard_list, get_storage_size)
220 # If we only have one shard, we return it => no need to build the index
221 if nb_shards == 1:
222 filename = filename_pattern.format(suffix="")
223 keys_in_shards = [k for k, t in state_dict.items() if not isinstance(t, str)]
224 return _StateDictSplit(
225 metadata={"total_size": total_size},
226 filename_to_tensors={filename: keys_in_shards},
227 tensor_to_filename={key: filename for key in keys_in_shards},
228 )
230 # Now that each tensor is assigned to a shard, let's assign a filename to each shard
231 tensor_name_to_filename = {}
232 filename_to_tensors = {}
233 for idx, shard in enumerate(shard_list):
234 filename = filename_pattern.format(suffix=f"-{idx + 1:05d}-of-{nb_shards:05d}")
235 for key in shard:
236 tensor_name_to_filename[key] = filename
237 filename_to_tensors[filename] = list(shard.keys())
239 # Build the index and return
240 return _StateDictSplit(
241 metadata={"total_size": total_size},
242 filename_to_tensors=filename_to_tensors,
243 tensor_to_filename=tensor_name_to_filename,
244 )
247def save_state_dict_as_huggingface_format(
248 save_directory: str | os.PathLike[str],
249 state_dict: dict[str, Any],
250 max_shard_size: str | int = "5GB",
251) -> None:
252 """Write ``state_dict`` to ``save_directory`` as Hugging Face-style safetensors (optionally sharded).
254 Args:
255 save_directory: Output directory path.
256 state_dict: Parameter tensors keyed by name.
257 max_shard_size: Per-shard byte cap (int) or string such as ``\"5GB\"``.
259 Returns:
260 None
261 """
262 if state_dict is None:
263 raise ValueError("state_dict is None.")
265 _validate_state_dict(state_dict)
267 if os.path.isfile(save_directory):
268 raise ValueError(f"The save_directory {save_directory} should be a directory, but a file.")
270 os.makedirs(save_directory, exist_ok=True)
272 weights_name = _SAFE_WEIGHTS_NAME
273 state_dict_split = _split_state_dict_into_shards(state_dict, max_shard_size=max_shard_size)
275 index = None
276 if state_dict_split.is_sharded:
277 index = {
278 "metadata": state_dict_split.metadata,
279 "weight_map": state_dict_split.tensor_to_filename,
280 }
282 filename_to_tensors = state_dict_split.filename_to_tensors.items()
284 for shard_file, tensors in filename_to_tensors:
285 shard = {tensor: state_dict[tensor] for tensor in tensors}
286 save_checkpoint_file(shard, os.path.join(save_directory, shard_file), ckpt_format="safetensors")
287 if index is None:
288 path_to_weights = os.path.join(save_directory, weights_name)
289 logger.info("Model weights saved in %s", path_to_weights)
290 else:
291 save_index_file = os.path.join(save_directory, _SAFE_WEIGHTS_INDEX_NAME)
292 with open(save_index_file, "w", encoding="utf-8") as f:
293 content = json.dumps(index, indent=2, sort_keys=True) + "\n"
294 f.write(content)
295 logger.info(
296 "Model exceeds max shard size %s; split into %s shards. Weight map index: %s",
297 max_shard_size,
298 len(state_dict_split.filename_to_tensors),
299 save_index_file,
300 )
303def parse_checkpoint_from_huggingface(resume_from_checkpoint: str | os.PathLike[str]) -> dict[str, Any]:
304 """Load a Hugging Face ``model.safetensors`` or sharded safetensors + index from a directory.
306 Args:
307 resume_from_checkpoint: Directory containing ``model.safetensors`` or index + shard files.
309 Returns:
310 Flattened str-keyed state dict (tensors on CPU when applicable).
311 """
312 safe_weights_file = os.path.join(resume_from_checkpoint, _SAFE_WEIGHTS_NAME)
314 state_dict: dict[str, Any] = {}
316 if os.path.isfile(safe_weights_file):
317 state_dict = load_checkpoint_file(safe_weights_file, ckpt_format="safetensors")
318 logger.info("Loaded safetensors checkpoint from %s", safe_weights_file)
320 else:
321 safe_index_file = os.path.join(resume_from_checkpoint, _SAFE_WEIGHTS_INDEX_NAME)
322 if not os.path.isfile(safe_index_file):
323 raise ValueError(f"Can't find a checkpoint index in {os.path.abspath(resume_from_checkpoint)}.")
325 with open(safe_index_file, "r", encoding="utf-8") as f:
326 index = json.load(f)
327 shard_files = list(set(index["weight_map"].values()))
329 total_size = 0
330 for shard_file in shard_files:
331 shard_path = os.path.join(resume_from_checkpoint, shard_file)
332 shard_state_dict = load_checkpoint_file(shard_path, ckpt_format="safetensors")
333 for key, value in shard_state_dict.items():
334 if key in state_dict:
335 logger.warning(
336 "Duplicate key %r when merging Hugging Face shards; keeping first occurrence.",
337 key,
338 )
339 continue
340 state_dict[key] = value
341 if isinstance(value, torch.Tensor):
342 size = get_tensor_storage_size(value)
343 total_size += size
344 logger.debug("Loaded tensor %r, size %s bytes", key, size)
345 else:
346 logger.debug("Loaded non-tensor entry %r", key)
347 logger.info("Merged Hugging Face shards, total tensor bytes (sum per key): %s", total_size)
348 _validate_state_dict(state_dict)
349 return state_dict
352def _full_checkpoint_format_for_path(path: str) -> str:
353 """Infer the :func:`load_checkpoint_file` ``ckpt_format`` from a full-checkpoint file suffix."""
354 lower = path.lower()
355 if lower.endswith(".safetensors"):
356 return "safetensors"
357 return "pickle"
360def full_state_dict_to_dcp_format(
361 state_dict: dict[str, Any],
362 dst_dir: str | os.PathLike[str],
363) -> None:
364 """Convert a full-weights ``state_dict`` into HyperParallel DCP layout under ``dst_dir``.
366 Args:
367 state_dict: Merged weights to store (see project conventions for the accepted value types).
368 dst_dir: Output directory for DCP files (``.metadata`` and shard files).
370 Returns:
371 None
372 """
373 _validate_state_dict(state_dict)
374 storage_writer = FileSystemWriter(dst_dir)
375 planner = StandardSavePlanner()
377 planner.configure_planner(state_dict=state_dict, is_coordinator=True)
378 storage_writer.configure_writer(is_coordinator=True, rank=0, use_collectives=True)
379 local_plan = planner.build_local_plan()
380 local_data = storage_writer.optimize_local_plan(local_plan)
382 central_plan, global_metadata = planner.build_global_plan([local_data])
383 central_plan = storage_writer.optimize_global_plan(central_plan)
385 final_local_plan = planner.finalize_plan(central_plan)
386 all_writes = storage_writer.execute_write(final_local_plan, planner)
388 storage_writer.finalize_checkpoint(metadata=global_metadata, results=[all_writes])
391def dcp_to_full_state_dict(src_dir: str | os.PathLike[str]) -> dict[str, Any]:
392 """Load a **shard** DCP layout from ``src_dir`` into one merged full-weights ``state_dict``.
394 Coordinator / single-process merge of the DCP shards on disk.
396 Args:
397 src_dir: Root directory of a saved DCP checkpoint.
399 Returns:
400 Merged state dictionary populated by the load planner and storage reader.
401 """
402 state_dict: dict[str, Any] = {}
403 planner = _DcpMergeLoadPlanner()
404 storage_reader = FileSystemReader(src_dir)
406 metadata = storage_reader.load_metadata()
407 planner.configure_planner(state_dict, metadata, is_coordinator=True)
408 storage_reader.configure_reader(metadata, is_coordinator=True, rank=0)
410 local_plan = planner.build_local_plan()
411 local_data = storage_reader.optimize_local_plan(local_plan)
413 # build_global_plan decides across ranks who reads which shard, so it takes every rank's
414 # plan and hands back the one belonging to the rank running it. This reads a checkpoint
415 # on its own, so the only plan among them is its own.
416 central_plan = planner.build_global_plan([local_data])
417 central_plan = storage_reader.optimize_global_plan(central_plan)
419 final_local_plan = planner.finalize_plan(central_plan)
420 storage_reader.execute_read(final_local_plan, planner)
422 return state_dict
425def convert_full_checkpoint_to_dcp(
426 src_ckpt: str | os.PathLike[str],
427 dst_dir: str | os.PathLike[str],
428 src_platform: Literal["torch", "huggingface"] = "torch",
429) -> None:
430 """Load checkpoint weights and write HyperParallel DCP under ``dst_dir``.
432 Args:
433 src_ckpt: Source path: Hugging Face safetensors directory (full or shard layout), or a PyTorch
434 **full checkpoint** file (pickle or ``.safetensors``).
435 dst_dir: Output directory for DCP files.
436 src_platform: ``huggingface`` (HF directory) or ``torch`` (full-checkpoint file read by
437 :func:`load_checkpoint_file`).
439 Returns:
440 None
441 """
442 if src_platform == "huggingface":
443 state_dict = parse_checkpoint_from_huggingface(src_ckpt)
444 elif src_platform == "torch":
445 if not os.path.isfile(src_ckpt):
446 raise ValueError(
447 "src_platform='torch' requires a single PyTorch checkpoint file path; "
448 f"src_ckpt must be an existing file, got {src_ckpt!r}."
449 )
450 fmt = _full_checkpoint_format_for_path(str(src_ckpt))
451 state_dict = load_checkpoint_file(str(src_ckpt), ckpt_format=fmt)
452 else:
453 raise ValueError(
454 f"Unsupported src_platform={src_platform!r}; expected 'huggingface' or 'torch'."
455 )
456 full_state_dict_to_dcp_format(state_dict, dst_dir)