Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / checkpoint_io.py: 60%

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"""Reading and writing one checkpoint file. 

16 

17The whole-file end of checkpoint storage: a flat tensor dict in, a file out, and back. 

18Which shards go in which file is decided elsewhere -- see ``filesystem_storage`` for the 

19save and load that call these. 

20""" 

21import torch 

22from safetensors.torch import load_file, save_file 

23 

24 

25def save_checkpoint_file(state_dict: dict, file_path: str, ckpt_format: str = "safetensors") -> None: 

26 """ 

27 Write a flat tensor dict to one checkpoint file. 

28 

29 Args: 

30 state_dict (dict): Tensors to write, keyed by their physical name. 

31 file_path (str): Where to write them. 

32 ckpt_format (str): ``safetensors`` for the safetensors format, anything else for 

33 a pickle written by :func:`torch.save`. Default ``safetensors``. 

34 """ 

35 if ckpt_format == "safetensors": 

36 save_file(tensors=state_dict, filename=file_path) 

37 else: 

38 torch.save(obj=state_dict, f=file_path) 

39 

40 

41def load_checkpoint_file(file_path: str, ckpt_format: str = "safetensors") -> dict: 

42 """ 

43 Read one checkpoint file back into a flat tensor dict. 

44 

45 Args: 

46 file_path (str): The file to read. 

47 ckpt_format (str): ``safetensors`` for the safetensors format, anything else for 

48 a pickle read by :func:`torch.load`. Default ``safetensors``. 

49 

50 Returns: 

51 dict: What the file holds, keyed by physical name. 

52 """ 

53 if ckpt_format == "safetensors": 

54 return load_file(filename=file_path) 

55 return torch.load(f=file_path, weights_only=True)