Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / fully_shard / utils.py: 94%

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"""Common policy and mesh metadata for fully_shard APIs.""" 

16from dataclasses import dataclass 

17from typing import Optional 

18 

19from hyper_parallel.collectives.cc import get_group_local_rank 

20from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

21from hyper_parallel.platform import get_platform 

22 

23platform = get_platform() 

24 

25 

26@dataclass 

27class MixedPrecisionPolicy: 

28 """ 

29 Configures mixed precision training for HSDP. 

30 

31 This policy controls data type casting during forward/backward computation 

32 and gradient reduction, enabling memory savings and potential speedups. 

33 

34 Attributes: 

35 param_dtype: Data type for parameter computation. If None, uses original dtype. 

36 reduce_dtype: Data type for gradient reduction. If None, uses param_dtype. 

37 output_dtype: Data type for module outputs. If None, no casting applied. 

38 """ 

39 param_dtype: Optional[platform.dtype] = None 

40 reduce_dtype: Optional[platform.dtype] = None 

41 output_dtype: Optional[platform.dtype] = None 

42 cast_forward_inputs: bool = True 

43 apply_grad_on_fp32_main_grad: bool = False 

44 

45 

46@dataclass 

47class OffloadPolicy: 

48 """ 

49 Base class for offload policies. 

50 

51 This represents no offloading and serves as the default policy. 

52 Subclass this to implement custom offload strategies. 

53 """ 

54 

55 

56@dataclass 

57class CPUOffloadPolicy(OffloadPolicy): 

58 """ 

59 Offloads sharded parameters and gradients to CPU memory. 

60 

61 When enabled, sharded parameters are kept on CPU and copied to device 

62 before all-gather. Gradients are copied back to CPU after backward. 

63 This reduces NPU memory usage at the cost of additional data transfers. 

64 

65 Attributes: 

66 pin_memory: If True, pins CPU memory for faster H2D/D2H transfers 

67 and enables overlap with computation. Disable if CPU memory 

68 is constrained. (Default: True) 

69 """ 

70 pin_memory: bool = True 

71 

72 

73@dataclass 

74class DataParallelMeshInfo: 

75 mesh: DeviceMesh 

76 shard_mesh_dim: Optional[int] = None 

77 replicate_mesh_dim: Optional[int] = None 

78 

79 def __post_init__(self): 

80 if self.shard_mesh_dim is None and self.replicate_mesh_dim is None: 

81 raise AssertionError( 

82 "At least one of shard_mesh_dim and replicate_mesh_dim must not be None" 

83 ) 

84 

85 

86@dataclass 

87class FSDPMeshInfo(DataParallelMeshInfo): 

88 def __post_init__(self): 

89 super().__post_init__() 

90 if self.shard_mesh_dim is None: 

91 raise AssertionError("Expects non-None shard_mesh_dim") 

92 self.shard_mesh_size: int = self.mesh.mesh_shape[self.shard_mesh_dim] 

93 self.shard_process_group = self.mesh.get_group(self.shard_mesh_dim) 

94 self.shard_mesh_rank: int = get_group_local_rank(self.shard_process_group) 

95 

96 

97@dataclass 

98class DDPMeshInfo(DataParallelMeshInfo): 

99 def __post_init__(self): 

100 super().__post_init__() 

101 if self.replicate_mesh_dim is None: 

102 raise AssertionError("Expects non-None replicate_mesh_dim") 

103 self.replicate_mesh_size: int = self.mesh.mesh_shape[self.replicate_mesh_dim] 

104 self.replicate_process_group = self.mesh.get_group(self.replicate_mesh_dim) 

105 self.replicate_mesh_rank: int = get_group_local_rank(self.replicate_process_group) 

106 

107 

108@dataclass 

109class HSDPMeshInfo(FSDPMeshInfo, DDPMeshInfo): 

110 # pylint: disable=W0246 

111 def __post_init__(self): 

112 # Calls `FSDPMeshInfo` -> `DDPMeshInfo` -> `DataParallelMeshInfo` 

113 super().__post_init__()