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

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.core.dtensor.placement_types import Placement 

22from hyper_parallel.platform import get_platform 

23 

24platform = get_platform() 

25DType = platform.dtype 

26 

27@dataclass 

28class MixedPrecisionPolicy: 

29 """ 

30 Configures mixed precision training for HSDP. 

31 

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

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

34 

35 Attributes: 

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

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

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

39 """ 

40 param_dtype: Optional[DType] = None 

41 reduce_dtype: Optional[DType] = None 

42 output_dtype: Optional[DType] = None 

43 cast_forward_inputs: bool = True 

44 apply_grad_on_fp32_main_grad: bool = False 

45 

46 

47@dataclass 

48class OffloadPolicy: 

49 """ 

50 Base class for offload policies. 

51 

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

53 Subclass this to implement custom offload strategies. 

54 """ 

55 

56 

57@dataclass 

58class CPUOffloadPolicy(OffloadPolicy): 

59 """ 

60 Offloads sharded parameters and gradients to CPU memory. 

61 

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

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

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

65 

66 Attributes: 

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

68 and enables overlap with computation. Disable if CPU memory 

69 is constrained. (Default: True) 

70 """ 

71 pin_memory: bool = True 

72 

73@dataclass 

74class CommFusionPolicy(): 

75 enable_comm_fusion: bool = False 

76 comm_fusion_zero_copy: bool = False 

77 

78 

79@dataclass 

80class DataParallelMeshInfo: 

81 mesh: DeviceMesh 

82 shard_mesh_dim: Optional[int] = None 

83 replicate_mesh_dim: Optional[int] = None 

84 

85 def __post_init__(self): 

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

87 raise AssertionError( 

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

89 ) 

90 

91 

92@dataclass 

93class FSDPMeshInfo(DataParallelMeshInfo): 

94 def __post_init__(self): 

95 super().__post_init__() 

96 if self.shard_mesh_dim is None: 

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

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

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

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

101 

102 

103@dataclass 

104class DDPMeshInfo(DataParallelMeshInfo): 

105 def __post_init__(self): 

106 super().__post_init__() 

107 if self.replicate_mesh_dim is None: 

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

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

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

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

112 

113 

114@dataclass 

115class HSDPMeshInfo(FSDPMeshInfo, DDPMeshInfo): 

116 # pylint: disable=W0246 

117 def __post_init__(self): 

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

119 super().__post_init__() 

120 

121 

122@dataclass(frozen=True) 

123class SourceShardMetaInfo: 

124 """Describe a parameter's source TP/EP layout before fully_shard.""" 

125 

126 mesh: DeviceMesh 

127 placements: tuple[Placement, ...] 

128 origin_is_dtensor: bool = False