Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / hf_storage.py: 98%
122 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"""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 "F16": "torch.float16",
84 "BF16": "torch.bfloat16",
85 "F32": "torch.float32",
86 "F64": "torch.float64",
87 "C64": "torch.complex64",
88}
91def _loads(document: Any, source: str) -> Any:
92 """
93 Parse a JSON document, saying which one it was when it does not parse.
95 Args:
96 document (Any): The text or bytes to parse.
97 source (str): What the document is, for the error.
99 Returns:
100 Any: The parsed document.
102 Raises:
103 ValueError: If the document is not JSON.
104 """
105 try:
106 return json.loads(document)
107 except (TypeError, ValueError) as e:
108 raise ValueError(f"{source} is not valid JSON: {e}") from e
111def _read_header(path: Path) -> dict[str, Any]:
112 """
113 Read the JSON header a safetensors file opens with, and nothing past it.
115 The header is an eight-byte little-endian length followed by that many bytes of JSON.
117 Args:
118 path (Path): The safetensors file.
120 Returns:
121 dict[str, Any]: One entry per tensor, plus ``__metadata__`` when the writer left one.
123 Raises:
124 ValueError: If the file is too short for the header it declares, or the header is not a
125 JSON object.
126 """
127 with open(path, "rb") as f:
128 file_size = os.fstat(f.fileno()).st_size
129 prefix = f.read(_HEADER_LENGTH_BYTES)
130 header_length = struct.unpack("<Q", prefix)[0] if len(prefix) == _HEADER_LENGTH_BYTES else None
131 if header_length is None or header_length > file_size - _HEADER_LENGTH_BYTES:
132 raise ValueError(f"{path} is not a safetensors file: it is too short for its header.")
133 header = _loads(f.read(header_length), f"The header of {path}")
134 if not isinstance(header, dict):
135 raise ValueError(f"The header of {path} is not a JSON object.")
136 return header
139def _sharding_info(header: dict[str, Any], path: Path) -> Optional[dict[str, Any]]:
140 """
141 The offsets torch's Hugging Face writer recorded for the tensors of one file, if it did.
143 Args:
144 header (dict[str, Any]): The header of the file.
145 path (Path): The file, for the error.
147 Returns:
148 Optional[dict[str, Any]]: ``{tensor name: {"saved_offsets": [...]}}``, or None for a file
149 that records none, which is every file Hugging Face writes itself.
151 Raises:
152 ValueError: If the recorded offsets are not a JSON object.
153 """
154 extra = header.get(_HEADER_METADATA_KEY)
155 if not isinstance(extra, dict) or DCP_SHARDING_INFO_KEY not in extra:
156 return None
157 info = _loads(extra[DCP_SHARDING_INFO_KEY], f"{DCP_SHARDING_INFO_KEY} of {path}")
158 if not isinstance(info, dict):
159 raise ValueError(f"{DCP_SHARDING_INFO_KEY} of {path} is not a JSON object.")
160 return info
163def _dtype_and_shape(entry: Any, key: str, path: Path) -> tuple[str, tuple[int, ...]]:
164 """
165 The dtype, spelled as DCP spells it, and the shape of one tensor entry of a header.
167 Args:
168 entry (Any): The header entry of the tensor.
169 key (str): Name of the tensor, for the error.
170 path (Path): The file, for the error.
172 Returns:
173 tuple[str, tuple[int, ...]]: The dtype and the shape.
175 Raises:
176 ValueError: If the entry has no valid shape, or a dtype safetensors does not read into torch.
177 """
178 shape = entry.get("shape") if isinstance(entry, dict) else None
179 if not isinstance(shape, list) or not all(isinstance(dim, int) and dim >= 0 for dim in shape):
180 raise ValueError(f"Tensor {key!r} in {path} has no valid shape: {entry!r}.")
181 code = entry.get("dtype")
182 if not isinstance(code, str) or code not in _SAFETENSORS_DTYPES:
183 raise ValueError(
184 f"Tensor {key!r} in {path} has dtype {code!r}, which is none of {', '.join(_SAFETENSORS_DTYPES)}."
185 )
186 return _SAFETENSORS_DTYPES[code], tuple(shape)
189def _saved_offsets(
190 sharding_info: Optional[dict[str, Any]], key: str, ndim: int, path: Path
191) -> tuple[int, ...]:
192 """
193 Where in the whole tensor the part of it stored under ``key`` in one file starts.
195 Args:
196 sharding_info (Optional[dict[str, Any]]): The offsets the file records, as
197 :func:`_sharding_info` returns them.
198 key (str): Name of the tensor.
199 ndim (int): Number of dimensions of the stored part.
200 path (Path): The file, for the error.
202 Returns:
203 tuple[int, ...]: The offsets, all zero for a file that records none.
205 Raises:
206 ValueError: If the file records offsets but not a valid set for this tensor. Taking it to
207 start at the origin instead could lay it over another shard of the same tensor.
208 """
209 if sharding_info is None:
210 return (0,) * ndim
211 entry = sharding_info.get(key)
212 offsets = entry.get(SAVED_OFFSETS_KEY) if isinstance(entry, dict) else None
213 if (not isinstance(offsets, list) or len(offsets) != ndim
214 or not all(isinstance(offset, int) and offset >= 0 for offset in offsets)):
215 raise ValueError(
216 f"{DCP_SHARDING_INFO_KEY} of {path} has no valid {SAVED_OFFSETS_KEY} for the "
217 f"{ndim}-dimensional tensor {key!r}: {entry!r}."
218 )
219 return tuple(offsets)
222def _read_chunks(path: Path, sharded: bool) -> Iterator[tuple[str, str, ChunkStorageMetadata]]:
223 """
224 Every tensor one safetensors file holds, as a chunk of the tensor its name stands for.
226 Args:
227 path (Path): The safetensors file.
228 sharded (bool): Whether the file is a shard written by torch, which has to record where
229 each of its tensors starts. Otherwise every tensor in it has to be whole.
231 Yields:
232 tuple[str, str, ChunkStorageMetadata]: Name, dtype, and where in the whole tensor it sits.
234 Raises:
235 ValueError: If a shard records no offsets, or a file of whole tensors records a tensor as
236 starting anywhere but the origin.
237 """
238 header = _read_header(path)
239 sharding_info = _sharding_info(header, path)
240 if sharded and sharding_info is None:
241 raise ValueError(
242 f"{path} records no {DCP_SHARDING_INFO_KEY}, so it is not a shard written by torch's "
243 f"HuggingFaceStorageWriter; read a Hugging Face checkpoint with HuggingFaceStorageReader."
244 )
245 for key, entry in header.items():
246 if key == _HEADER_METADATA_KEY:
247 continue
248 dtype, shape = _dtype_and_shape(entry, key, path)
249 offsets = _saved_offsets(sharding_info, key, len(shape), path)
250 if not sharded and any(offsets):
251 raise ValueError(
252 f"Tensor {key!r} in {path} is the shard at offsets {offsets} of a larger tensor; read a "
253 f"directory of shards written by torch with TorchShardedSafetensorsReader."
254 )
255 yield key, dtype, ChunkStorageMetadata(offsets=offsets, sizes=shape)
258def _tensor_size(key: str, chunks: list[ChunkStorageMetadata]) -> tuple[int, ...]:
259 """
260 The size of the tensor a set of chunks makes up, checking that they make up exactly that.
262 Args:
263 key (str): Name of the tensor, for the error.
264 chunks (list[ChunkStorageMetadata]): Every chunk of it, no two at the same offsets.
266 Returns:
267 tuple[int, ...]: The smallest size that holds every chunk.
269 Raises:
270 ValueError: If the chunks differ in number of dimensions, or hold more or fewer elements
271 than that size does: a file of the checkpoint is missing, or two chunks overlap.
272 """
273 ndim = len(chunks[0].sizes)
274 if any(len(chunk.sizes) != ndim for chunk in chunks):
275 raise ValueError(f"The chunks of tensor {key!r} differ in number of dimensions.")
276 size = tuple(max(chunk.offsets[dim] + chunk.sizes[dim] for chunk in chunks) for dim in range(ndim))
277 held = sum(prod(chunk.sizes) for chunk in chunks)
278 if held != prod(size):
279 raise ValueError(
280 f"The chunks of tensor {key!r} hold {held} elements between them, where a tensor of size "
281 f"{size} has {prod(size)}: a file of the checkpoint is missing, or two chunks overlap."
282 )
283 return size
286def _describe_checkpoint(checkpoint_dir: Path, file_names: list[str], sharded: bool) -> Metadata:
287 """
288 Build the metadata of a checkpoint out of the headers of its safetensors files.
290 Args:
291 checkpoint_dir (Path): The checkpoint directory.
292 file_names (list[str]): The files of the checkpoint, relative to ``checkpoint_dir``.
293 sharded (bool): Whether the files are shards written by torch; see :func:`_read_chunks`.
295 Returns:
296 Metadata: Size, dtype and chunks of every tensor, and the file each chunk is in.
298 Raises:
299 ValueError: If a file is not of the kind ``sharded`` says, a header is malformed, or the
300 chunks of one tensor do not fit together: the same chunk twice, chunks in different
301 dtypes, or chunks that do not make up the whole tensor.
302 """
303 chunks: dict[str, list[ChunkStorageMetadata]] = {}
304 dtypes: dict[str, str] = {}
305 storage_data: dict[MetadataIndex, StorageInfo] = {}
306 for file_name in file_names:
307 for key, dtype, chunk in _read_chunks(checkpoint_dir / file_name, sharded):
308 index = MetadataIndex(fqn=key, offset=chunk.offsets)
309 if index in storage_data:
310 raise ValueError(
311 f"Tensor {key!r} at offsets {chunk.offsets} is in both "
312 f"{storage_data[index].relative_path} and {file_name}."
313 )
314 if dtypes.setdefault(key, dtype) != dtype:
315 raise ValueError(f"Tensor {key!r} is {dtype} in {file_name} but {dtypes[key]} in another file.")
316 chunks.setdefault(key, []).append(chunk)
317 storage_data[index] = StorageInfo(relative_path=file_name, offset=0, length=-1, tensor_key=key)
318 state_dict_metadata = {
319 key: TensorStorageMetadata(
320 properties=TensorProperties(dtype=dtypes[key]),
321 size=_tensor_size(key, key_chunks),
322 chunks=key_chunks,
323 )
324 for key, key_chunks in chunks.items()
325 }
326 return Metadata(state_dict_metadata=state_dict_metadata, storage_data=storage_data)
329class HuggingFaceStorageReader(FileSystemReader):
330 """
331 Storage reader for a Hugging Face checkpoint: whole tensors, each in exactly one file.
333 Pass it to :func:`load` in place of a checkpoint path::
335 load(model.state_dict(), storage_reader=HuggingFaceStorageReader("/models/Qwen3-8B"))
337 The files read are the ones ``model.safetensors.index.json`` names, or ``model.safetensors``
338 alone when there is no index. Nothing else in the directory is read, such as a second copy of
339 the weights in a file of another name. A file whose header records offsets the way torch's
340 HuggingFaceStorageWriter does is read as long as every offset is zero: anything else is a shard
341 of a larger tensor, and a directory of those is for :class:`TorchShardedSafetensorsReader`.
343 The state dict is filled by name, as from any DCP checkpoint: each of its keys has to name a
344 tensor in the files, at the same global shape. Beyond that the target is free. A DTensor may
345 be laid out any way, and each rank reads only its own slice of the file. The dtype may differ
346 from the file's, and the copy into the target converts it. Ranks holding the same slice read
347 it once and broadcast it, as :func:`load` does for any checkpoint.
349 A key the files do not have fails the load, unless the planner is
350 ``StandardLoadPlanner(allow_partial_load=True)``. A model with tied weights needs that, since
351 Hugging Face stores a tied tensor once. Quantized tensors are read as they are stored: an FP8
352 weight arrives as FP8 values, and its scales are tensors of their own.
353 """
355 @dcp_timer_decorator
356 def load_metadata(self, **kwargs: Any) -> Metadata:
357 """
358 Describe the checkpoint from the headers of its safetensors files, and nothing past them.
360 Args:
361 **kwargs (Any): Ignored. :func:`load` asks again with ``rank`` when a first ask finds
362 nothing, and a Hugging Face checkpoint has no rank-local metadata to find instead.
364 Returns:
365 Metadata: Size, dtype and file of every tensor.
367 Raises:
368 FileNotFoundError: If the directory holds neither the index nor the single weights
369 file, or a file the index names is not there.
370 ValueError: If the index or a header is malformed, a dtype has no torch counterpart, a
371 tensor is a shard of a larger one, or two files hold the same tensor.
372 """
373 return _describe_checkpoint(self.checkpoint_dir, self._weight_files(), sharded=False)
375 def _weight_files(self) -> list[str]:
376 """
377 Name the files that make up the checkpoint.
379 Returns:
380 list[str]: File names relative to the checkpoint directory, sorted.
382 Raises:
383 FileNotFoundError: If the directory holds neither the index nor the single weights file.
384 ValueError: If the index does not map tensor names to file names.
385 """
386 index_path = self.checkpoint_dir / _SAFE_WEIGHTS_INDEX_NAME
387 if index_path.is_file():
388 index = _loads(index_path.read_bytes(), str(index_path))
389 weight_map = index.get("weight_map") if isinstance(index, dict) else None
390 if (not isinstance(weight_map, dict) or not weight_map
391 or not all(isinstance(file_name, str) for file_name in weight_map.values())):
392 raise ValueError(f"{index_path} has no weight_map mapping tensor names to file names.")
393 return sorted(set(weight_map.values()))
394 if not (self.checkpoint_dir / _SAFE_WEIGHTS_NAME).is_file():
395 raise FileNotFoundError(
396 f"Neither {_SAFE_WEIGHTS_INDEX_NAME} nor {_SAFE_WEIGHTS_NAME} found in {self.checkpoint_dir}. "
397 f"A directory of shards written by torch is read with TorchShardedSafetensorsReader."
398 )
399 return [_SAFE_WEIGHTS_NAME]
402class TorchShardedSafetensorsReader(FileSystemReader):
403 """
404 Storage reader for the shards torch's ``HuggingFaceStorageWriter`` writes with ``save_distributed``.
406 Saving that way without consolidating, every rank writes its shard of each tensor to a file of
407 its own, ``shard-00001-model-00001-of-00001.safetensors`` and so on, under the tensor's own
408 name, and records in the file's header where in the whole tensor the shard sits. No index is
409 written, since a tensor is spread over several files. torchtitan writes such a directory as the
410 ``sharded`` directory of a checkpoint it saves in Hugging Face format, before consolidating it.
412 Every safetensors file in the directory is read, and every one has to record those offsets.
413 The shards of a tensor are its chunks, and have to make up the whole tensor between them,
414 leaving no gap and overlapping nowhere. A directory with an index holds whole tensors, and is
415 for :class:`HuggingFaceStorageReader`.
417 Loading from it is as from any DCP checkpoint; see :class:`HuggingFaceStorageReader` for what
418 the state dict may be.
419 """
421 @dcp_timer_decorator
422 def load_metadata(self, **kwargs: Any) -> Metadata:
423 """
424 Describe the checkpoint from the headers of its shard files, and nothing past them.
426 Args:
427 **kwargs (Any): Ignored, as :meth:`HuggingFaceStorageReader.load_metadata` ignores them.
429 Returns:
430 Metadata: Size, dtype and chunks of every tensor, and the file each chunk is in.
432 Raises:
433 FileNotFoundError: If the directory holds no safetensors file.
434 ValueError: If the directory has an index, a file records no offsets, a header is
435 malformed, a dtype has no torch counterpart, or the shards of one tensor do not fit
436 together: the same shard twice, shards in different dtypes, or shards that do not
437 make up the whole tensor.
438 """
439 indexes = sorted(path.name for path in self.checkpoint_dir.glob(f"*{SAFETENSORS_INDEX_SUFFIX}"))
440 if indexes:
441 raise ValueError(
442 f"{self.checkpoint_dir} has {indexes[0]}, which a checkpoint of whole tensors has and one "
443 f"of shards does not; read it with HuggingFaceStorageReader."
444 )
445 files = sorted(path.name for path in self.checkpoint_dir.glob(f"*{SAFETENSORS_SUFFIX}"))
446 if not files:
447 raise FileNotFoundError(f"No *{SAFETENSORS_SUFFIX} file found in {self.checkpoint_dir}.")
448 return _describe_checkpoint(self.checkpoint_dir, files, sharded=True)