Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / checkpoint / hf_storage.py: 98%
122 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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"""Safetensors checkpoints written outside DCP, loaded through the DCP read path.
17Two kinds of directory are read, each by a reader of its own, since what makes a directory of one
18kind valid makes one of the other kind wrong:
20* :class:`HuggingFaceStorageReader` reads a Hugging Face checkpoint, as ``save_pretrained`` writes
21 one: whole tensors, each in exactly one file, which is ``model.safetensors`` or one of the files a
22 ``model.safetensors.index.json`` names.
23* :class:`TorchShardedSafetensorsReader` reads what torch's ``HuggingFaceStorageWriter`` writes with
24 ``save_distributed=True`` before consolidating it: every rank's shard of a tensor in a file of its
25 own, where the shard sits in the whole tensor recorded in the file's header, and no index.
27A safetensors file opens with a JSON header giving the dtype and shape of every tensor in it, and
28that is all a load plans against. Both readers build the checkpoint :class:`Metadata` out of those
29headers, where :class:`FileSystemReader` unpickles a ``.metadata`` file, and inherit the rest of the
30load: planning, resharding, the replicated-shard broadcast and the reads themselves.
32Tensors are known by the names the files give them. Mapping those onto the names of a model, and
33merging or splitting tensors laid out differently from it, is left to the layer above.
34"""
35import json
36import os
37import struct
38from collections.abc import Iterator
39from math import prod
40from pathlib import Path
41from typing import Any, Optional
43from hyper_parallel.core.distributed_checkpoint.filesystem_storage import FileSystemReader
44from hyper_parallel.core.distributed_checkpoint.metadata import (
45 ChunkStorageMetadata,
46 Metadata,
47 MetadataIndex,
48 TensorProperties,
49 TensorStorageMetadata,
50)
51from hyper_parallel.core.distributed_checkpoint.storage import StorageInfo
52from hyper_parallel.core.distributed_checkpoint.utils import dcp_timer_decorator
54SAFETENSORS_SUFFIX = ".safetensors"
55SAFETENSORS_INDEX_SUFFIX = ".safetensors.index.json"
56# torch's HuggingFaceStorageWriter records where in the whole tensor each tensor of a file starts: a
57# JSON string under this key of the header's ``__metadata__``, mapping every tensor name in the file
58# to ``{"saved_offsets": [...]}``. The format allows nothing but strings there, hence JSON in a string.
59DCP_SHARDING_INFO_KEY = "DCP_SHARDING_INFO"
60SAVED_OFFSETS_KEY = "saved_offsets"
62_SAFE_WEIGHTS_NAME = "model.safetensors"
63_SAFE_WEIGHTS_INDEX_NAME = "model.safetensors.index.json"
64_HEADER_METADATA_KEY = "__metadata__"
65_HEADER_LENGTH_BYTES = 8
67# Every dtype code safetensors (0.8) reads into torch, with the dtype it is read as, spelled the
68# way a DCP checkpoint records one: ``str(tensor.dtype)``.
69_SAFETENSORS_DTYPES = {
70 "BOOL": "torch.bool",
71 "U8": "torch.uint8",
72 "I8": "torch.int8",
73 "U16": "torch.uint16",
74 "I16": "torch.int16",
75 "U32": "torch.uint32",
76 "I32": "torch.int32",
77 "U64": "torch.uint64",
78 "I64": "torch.int64",
79 "F8_E4M3": "torch.float8_e4m3fn",
80 "F8_E4M3FNUZ": "torch.float8_e4m3fnuz",
81 "F8_E5M2": "torch.float8_e5m2",
82 "F8_E5M2FNUZ": "torch.float8_e5m2fnuz",
83 "F8_E8M0": "torch.float8_e8m0fnu",
84 "F16": "torch.float16",
85 "BF16": "torch.bfloat16",
86 "F32": "torch.float32",
87 "F64": "torch.float64",
88 "C64": "torch.complex64",
89}
92def _loads(document: Any, source: str) -> Any:
93 """
94 Parse a JSON document, saying which one it was when it does not parse.
96 Args:
97 document (Any): The text or bytes to parse.
98 source (str): What the document is, for the error.
100 Returns:
101 Any: The parsed document.
103 Raises:
104 ValueError: If the document is not JSON.
105 """
106 try:
107 return json.loads(document)
108 except (TypeError, ValueError) as e:
109 raise ValueError(f"{source} is not valid JSON: {e}") from e
112def _read_header(path: Path) -> dict[str, Any]:
113 """
114 Read the JSON header a safetensors file opens with, and nothing past it.
116 The header is an eight-byte little-endian length followed by that many bytes of JSON.
118 Args:
119 path (Path): The safetensors file.
121 Returns:
122 dict[str, Any]: One entry per tensor, plus ``__metadata__`` when the writer left one.
124 Raises:
125 ValueError: If the file is too short for the header it declares, or the header is not a
126 JSON object.
127 """
128 with open(path, "rb") as f:
129 file_size = os.fstat(f.fileno()).st_size
130 prefix = f.read(_HEADER_LENGTH_BYTES)
131 header_length = struct.unpack("<Q", prefix)[0] if len(prefix) == _HEADER_LENGTH_BYTES else None
132 if header_length is None or header_length > file_size - _HEADER_LENGTH_BYTES:
133 raise ValueError(f"{path} is not a safetensors file: it is too short for its header.")
134 header = _loads(f.read(header_length), f"The header of {path}")
135 if not isinstance(header, dict):
136 raise ValueError(f"The header of {path} is not a JSON object.")
137 return header
140def _sharding_info(header: dict[str, Any], path: Path) -> Optional[dict[str, Any]]:
141 """
142 The offsets torch's Hugging Face writer recorded for the tensors of one file, if it did.
144 Args:
145 header (dict[str, Any]): The header of the file.
146 path (Path): The file, for the error.
148 Returns:
149 Optional[dict[str, Any]]: ``{tensor name: {"saved_offsets": [...]}}``, or None for a file
150 that records none, which is every file Hugging Face writes itself.
152 Raises:
153 ValueError: If the recorded offsets are not a JSON object.
154 """
155 extra = header.get(_HEADER_METADATA_KEY)
156 if not isinstance(extra, dict) or DCP_SHARDING_INFO_KEY not in extra:
157 return None
158 info = _loads(extra[DCP_SHARDING_INFO_KEY], f"{DCP_SHARDING_INFO_KEY} of {path}")
159 if not isinstance(info, dict):
160 raise ValueError(f"{DCP_SHARDING_INFO_KEY} of {path} is not a JSON object.")
161 return info
164def _dtype_and_shape(entry: Any, key: str, path: Path) -> tuple[str, tuple[int, ...]]:
165 """
166 The dtype, spelled as DCP spells it, and the shape of one tensor entry of a header.
168 Args:
169 entry (Any): The header entry of the tensor.
170 key (str): Name of the tensor, for the error.
171 path (Path): The file, for the error.
173 Returns:
174 tuple[str, tuple[int, ...]]: The dtype and the shape.
176 Raises:
177 ValueError: If the entry has no valid shape, or a dtype safetensors does not read into torch.
178 """
179 shape = entry.get("shape") if isinstance(entry, dict) else None
180 if not isinstance(shape, list) or not all(isinstance(dim, int) and dim >= 0 for dim in shape):
181 raise ValueError(f"Tensor {key!r} in {path} has no valid shape: {entry!r}.")
182 code = entry.get("dtype")
183 if not isinstance(code, str) or code not in _SAFETENSORS_DTYPES:
184 raise ValueError(
185 f"Tensor {key!r} in {path} has dtype {code!r}, which is none of {', '.join(_SAFETENSORS_DTYPES)}."
186 )
187 return _SAFETENSORS_DTYPES[code], tuple(shape)
190def _saved_offsets(
191 sharding_info: Optional[dict[str, Any]], key: str, ndim: int, path: Path
192) -> tuple[int, ...]:
193 """
194 Where in the whole tensor the part of it stored under ``key`` in one file starts.
196 Args:
197 sharding_info (Optional[dict[str, Any]]): The offsets the file records, as
198 :func:`_sharding_info` returns them.
199 key (str): Name of the tensor.
200 ndim (int): Number of dimensions of the stored part.
201 path (Path): The file, for the error.
203 Returns:
204 tuple[int, ...]: The offsets, all zero for a file that records none.
206 Raises:
207 ValueError: If the file records offsets but not a valid set for this tensor. Taking it to
208 start at the origin instead could lay it over another shard of the same tensor.
209 """
210 if sharding_info is None:
211 return (0,) * ndim
212 entry = sharding_info.get(key)
213 offsets = entry.get(SAVED_OFFSETS_KEY) if isinstance(entry, dict) else None
214 if (not isinstance(offsets, list) or len(offsets) != ndim
215 or not all(isinstance(offset, int) and offset >= 0 for offset in offsets)):
216 raise ValueError(
217 f"{DCP_SHARDING_INFO_KEY} of {path} has no valid {SAVED_OFFSETS_KEY} for the "
218 f"{ndim}-dimensional tensor {key!r}: {entry!r}."
219 )
220 return tuple(offsets)
223def _read_chunks(path: Path, sharded: bool) -> Iterator[tuple[str, str, ChunkStorageMetadata]]:
224 """
225 Every tensor one safetensors file holds, as a chunk of the tensor its name stands for.
227 Args:
228 path (Path): The safetensors file.
229 sharded (bool): Whether the file is a shard written by torch, which has to record where
230 each of its tensors starts. Otherwise every tensor in it has to be whole.
232 Yields:
233 tuple[str, str, ChunkStorageMetadata]: Name, dtype, and where in the whole tensor it sits.
235 Raises:
236 ValueError: If a shard records no offsets, or a file of whole tensors records a tensor as
237 starting anywhere but the origin.
238 """
239 header = _read_header(path)
240 sharding_info = _sharding_info(header, path)
241 if sharded and sharding_info is None:
242 raise ValueError(
243 f"{path} records no {DCP_SHARDING_INFO_KEY}, so it is not a shard written by torch's "
244 f"HuggingFaceStorageWriter; read a Hugging Face checkpoint with HuggingFaceStorageReader."
245 )
246 for key, entry in header.items():
247 if key == _HEADER_METADATA_KEY:
248 continue
249 dtype, shape = _dtype_and_shape(entry, key, path)
250 offsets = _saved_offsets(sharding_info, key, len(shape), path)
251 if not sharded and any(offsets):
252 raise ValueError(
253 f"Tensor {key!r} in {path} is the shard at offsets {offsets} of a larger tensor; read a "
254 f"directory of shards written by torch with TorchShardedSafetensorsReader."
255 )
256 yield key, dtype, ChunkStorageMetadata(offsets=offsets, sizes=shape)
259def _tensor_size(key: str, chunks: list[ChunkStorageMetadata]) -> tuple[int, ...]:
260 """
261 The size of the tensor a set of chunks makes up, checking that they make up exactly that.
263 Args:
264 key (str): Name of the tensor, for the error.
265 chunks (list[ChunkStorageMetadata]): Every chunk of it, no two at the same offsets.
267 Returns:
268 tuple[int, ...]: The smallest size that holds every chunk.
270 Raises:
271 ValueError: If the chunks differ in number of dimensions, or hold more or fewer elements
272 than that size does: a file of the checkpoint is missing, or two chunks overlap.
273 """
274 ndim = len(chunks[0].sizes)
275 if any(len(chunk.sizes) != ndim for chunk in chunks):
276 raise ValueError(f"The chunks of tensor {key!r} differ in number of dimensions.")
277 size = tuple(max(chunk.offsets[dim] + chunk.sizes[dim] for chunk in chunks) for dim in range(ndim))
278 held = sum(prod(chunk.sizes) for chunk in chunks)
279 if held != prod(size):
280 raise ValueError(
281 f"The chunks of tensor {key!r} hold {held} elements between them, where a tensor of size "
282 f"{size} has {prod(size)}: a file of the checkpoint is missing, or two chunks overlap."
283 )
284 return size
287def _describe_checkpoint(checkpoint_dir: Path, file_names: list[str], sharded: bool) -> Metadata:
288 """
289 Build the metadata of a checkpoint out of the headers of its safetensors files.
291 Args:
292 checkpoint_dir (Path): The checkpoint directory.
293 file_names (list[str]): The files of the checkpoint, relative to ``checkpoint_dir``.
294 sharded (bool): Whether the files are shards written by torch; see :func:`_read_chunks`.
296 Returns:
297 Metadata: Size, dtype and chunks of every tensor, and the file each chunk is in.
299 Raises:
300 ValueError: If a file is not of the kind ``sharded`` says, a header is malformed, or the
301 chunks of one tensor do not fit together: the same chunk twice, chunks in different
302 dtypes, or chunks that do not make up the whole tensor.
303 """
304 chunks: dict[str, list[ChunkStorageMetadata]] = {}
305 dtypes: dict[str, str] = {}
306 storage_data: dict[MetadataIndex, StorageInfo] = {}
307 for file_name in file_names:
308 for key, dtype, chunk in _read_chunks(checkpoint_dir / file_name, sharded):
309 index = MetadataIndex(fqn=key, offset=chunk.offsets)
310 if index in storage_data:
311 raise ValueError(
312 f"Tensor {key!r} at offsets {chunk.offsets} is in both "
313 f"{storage_data[index].relative_path} and {file_name}."
314 )
315 if dtypes.setdefault(key, dtype) != dtype:
316 raise ValueError(f"Tensor {key!r} is {dtype} in {file_name} but {dtypes[key]} in another file.")
317 chunks.setdefault(key, []).append(chunk)
318 storage_data[index] = StorageInfo(relative_path=file_name, offset=0, length=-1, tensor_key=key)
319 state_dict_metadata = {
320 key: TensorStorageMetadata(
321 properties=TensorProperties(dtype=dtypes[key]),
322 size=_tensor_size(key, key_chunks),
323 chunks=key_chunks,
324 )
325 for key, key_chunks in chunks.items()
326 }
327 return Metadata(state_dict_metadata=state_dict_metadata, storage_data=storage_data)
330class HuggingFaceStorageReader(FileSystemReader):
331 """
332 Storage reader for a Hugging Face checkpoint: whole tensors, each in exactly one file.
334 Pass it to :func:`load` in place of a checkpoint path::
336 load(model.state_dict(), storage_reader=HuggingFaceStorageReader("/models/Qwen3-8B"))
338 The files read are the ones ``model.safetensors.index.json`` names, or ``model.safetensors``
339 alone when there is no index. Nothing else in the directory is read, such as a second copy of
340 the weights in a file of another name. A file whose header records offsets the way torch's
341 HuggingFaceStorageWriter does is read as long as every offset is zero: anything else is a shard
342 of a larger tensor, and a directory of those is for :class:`TorchShardedSafetensorsReader`.
344 The state dict is filled by name, as from any DCP checkpoint: each of its keys has to name a
345 tensor in the files, at the same global shape. Beyond that the target is free. A DTensor may
346 be laid out any way, and each rank reads only its own slice of the file. The dtype may differ
347 from the file's, and the copy into the target converts it. Ranks holding the same slice read
348 it once and broadcast it, as :func:`load` does for any checkpoint.
350 A key the files do not have fails the load, unless the planner is
351 ``StandardLoadPlanner(allow_partial_load=True)``. A model with tied weights needs that, since
352 Hugging Face stores a tied tensor once. Quantized tensors are read as they are stored: an FP8
353 weight arrives as FP8 values, and its scales are tensors of their own.
354 """
356 @dcp_timer_decorator
357 def load_metadata(self, **kwargs: Any) -> Metadata:
358 """
359 Describe the checkpoint from the headers of its safetensors files, and nothing past them.
361 Args:
362 **kwargs (Any): Ignored. :func:`load` asks again with ``rank`` when a first ask finds
363 nothing, and a Hugging Face checkpoint has no rank-local metadata to find instead.
365 Returns:
366 Metadata: Size, dtype and file of every tensor.
368 Raises:
369 FileNotFoundError: If the directory holds neither the index nor the single weights
370 file, or a file the index names is not there.
371 ValueError: If the index or a header is malformed, a dtype has no torch counterpart, a
372 tensor is a shard of a larger one, or two files hold the same tensor.
373 """
374 return _describe_checkpoint(self.checkpoint_dir, self._weight_files(), sharded=False)
376 def _weight_files(self) -> list[str]:
377 """
378 Name the files that make up the checkpoint.
380 Returns:
381 list[str]: File names relative to the checkpoint directory, sorted.
383 Raises:
384 FileNotFoundError: If the directory holds neither the index nor the single weights file.
385 ValueError: If the index does not map tensor names to file names.
386 """
387 index_path = self.checkpoint_dir / _SAFE_WEIGHTS_INDEX_NAME
388 if index_path.is_file():
389 index = _loads(index_path.read_bytes(), str(index_path))
390 weight_map = index.get("weight_map") if isinstance(index, dict) else None
391 if (not isinstance(weight_map, dict) or not weight_map
392 or not all(isinstance(file_name, str) for file_name in weight_map.values())):
393 raise ValueError(f"{index_path} has no weight_map mapping tensor names to file names.")
394 return sorted(set(weight_map.values()))
395 if not (self.checkpoint_dir / _SAFE_WEIGHTS_NAME).is_file():
396 raise FileNotFoundError(
397 f"Neither {_SAFE_WEIGHTS_INDEX_NAME} nor {_SAFE_WEIGHTS_NAME} found in {self.checkpoint_dir}. "
398 f"A directory of shards written by torch is read with TorchShardedSafetensorsReader."
399 )
400 return [_SAFE_WEIGHTS_NAME]
403class TorchShardedSafetensorsReader(FileSystemReader):
404 """
405 Storage reader for the shards torch's ``HuggingFaceStorageWriter`` writes with ``save_distributed``.
407 Saving that way without consolidating, every rank writes its shard of each tensor to a file of
408 its own, ``shard-00001-model-00001-of-00001.safetensors`` and so on, under the tensor's own
409 name, and records in the file's header where in the whole tensor the shard sits. No index is
410 written, since a tensor is spread over several files. torchtitan writes such a directory as the
411 ``sharded`` directory of a checkpoint it saves in Hugging Face format, before consolidating it.
413 Every safetensors file in the directory is read, and every one has to record those offsets.
414 The shards of a tensor are its chunks, and have to make up the whole tensor between them,
415 leaving no gap and overlapping nowhere. A directory with an index holds whole tensors, and is
416 for :class:`HuggingFaceStorageReader`.
418 Loading from it is as from any DCP checkpoint; see :class:`HuggingFaceStorageReader` for what
419 the state dict may be.
420 """
422 @dcp_timer_decorator
423 def load_metadata(self, **kwargs: Any) -> Metadata:
424 """
425 Describe the checkpoint from the headers of its shard files, and nothing past them.
427 Args:
428 **kwargs (Any): Ignored, as :meth:`HuggingFaceStorageReader.load_metadata` ignores them.
430 Returns:
431 Metadata: Size, dtype and chunks of every tensor, and the file each chunk is in.
433 Raises:
434 FileNotFoundError: If the directory holds no safetensors file.
435 ValueError: If the directory has an index, a file records no offsets, a header is
436 malformed, a dtype has no torch counterpart, or the shards of one tensor do not fit
437 together: the same shard twice, shards in different dtypes, or shards that do not
438 make up the whole tensor.
439 """
440 indexes = sorted(path.name for path in self.checkpoint_dir.glob(f"*{SAFETENSORS_INDEX_SUFFIX}"))
441 if indexes:
442 raise ValueError(
443 f"{self.checkpoint_dir} has {indexes[0]}, which a checkpoint of whole tensors has and one "
444 f"of shards does not; read it with HuggingFaceStorageReader."
445 )
446 files = sorted(path.name for path in self.checkpoint_dir.glob(f"*{SAFETENSORS_SUFFIX}"))
447 if not files:
448 raise FileNotFoundError(f"No *{SAFETENSORS_SUFFIX} file found in {self.checkpoint_dir}.")
449 return _describe_checkpoint(self.checkpoint_dir, files, sharded=True)