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

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"""Optimizer state dict core API — a forwarding layer over the torch implementation. 

16 

17This module provides ``get_optim_state_dict`` and ``set_optim_state_dict`` as core-layer 

18entry points that forward to the FSDP-aware torch implementations, which know how to gather 

19DTensor shards and flatten an optimizer state dict. 

20""" 

21from __future__ import annotations 

22 

23from typing import Any 

24 

25from hyper_parallel.core.fully_shard.state_dict_utils import ( 

26 get_optim_state_dict as _get_optim_state_dict, 

27 set_optim_state_dict as _set_optim_state_dict, 

28) 

29 

30 

31def get_optim_state_dict(model: Any, optimizer: Any, *, options: Any = None) -> Any: 

32 """Get the optimizer state dict. 

33 

34 Args: 

35 model: The model whose parameters are optimized. 

36 optimizer: The optimizer instance. 

37 options: Optional configuration (full_state_dict, cpu_offload, 

38 flatten_optimizer_state_dict, etc.). 

39 

40 Returns: 

41 dict: Optimizer state dict with FQN-based keys. 

42 """ 

43 return _get_optim_state_dict(model, optimizer, options=options) 

44 

45 

46def set_optim_state_dict( 

47 model: Any, 

48 optimizer: Any, 

49 optim_state_dict: Any, 

50 *, 

51 options: Any = None, 

52) -> None: 

53 """Load an optimizer state dict into ``optimizer``. 

54 

55 Args: 

56 model: The model whose parameters are optimized. 

57 optimizer: The target optimizer instance. 

58 optim_state_dict: The optimizer state dict to load. 

59 options: Optional configuration (full_state_dict, cpu_offload, 

60 strict, broadcast_from_rank0, etc.). 

61 """ 

62 _set_optim_state_dict(model, optimizer, optim_state_dict, options=options)