Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / tensor_parallel / loss_parallel_ops_common.py: 39%

1# Copyright 2025-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 utilities for loss_parallel operations. 

16 

17This module contains PyTorch helper functions and validation logic. 

18""" 

19 

20from __future__ import annotations 

21 

22import warnings 

23from typing import Any, Optional, TYPE_CHECKING 

24 

25from hyper_parallel.core.dtensor.dtensor import DTensor 

26from hyper_parallel.core.dtensor.placement_types import Shard, Replicate 

27from hyper_parallel.core.tensor_parallel.loss_parallel import is_loss_parallel_active 

28 

29if TYPE_CHECKING: 

30 from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

31 

32__all__ = [ 

33 "_is_dtensor", 

34 "_is_shard_on_last_dim", 

35 "_is_replicate", 

36 "_get_mesh_and_dim", 

37 "_get_local_tensor", 

38 "_get_full_tensor", 

39 "_validate_target_type_base", 

40 "_validate_mesh_and_shard", 

41 "_validate_cross_entropy_params", 

42 "_check_context_and_layout", 

43] 

44 

45 

46def _is_dtensor(obj: Any) -> bool: 

47 """Check if object is a DTensor.""" 

48 return isinstance(obj, DTensor) 

49 

50 

51def _is_shard_on_last_dim(dtensor: DTensor) -> bool: 

52 """Check if DTensor is Shard on the last dimension.""" 

53 if not _is_dtensor(dtensor): 

54 return False 

55 placements = dtensor.placements 

56 if not placements: 

57 return False 

58 last_placement = placements[-1] 

59 if isinstance(last_placement, Shard): 

60 return last_placement.dim in (-1, len(dtensor.shape) - 1) 

61 return False 

62 

63 

64def _is_replicate(dtensor: DTensor) -> bool: 

65 """Check if DTensor is Replicate.""" 

66 if not _is_dtensor(dtensor): 

67 return False 

68 return all(isinstance(p, Replicate) for p in dtensor.placements) 

69 

70 

71def _get_mesh_and_dim(dtensor: DTensor) -> tuple[Optional["DeviceMesh"], Optional[int]]: 

72 """Get mesh and shard dimension from DTensor.""" 

73 mesh = dtensor.device_mesh 

74 shard_dim = None 

75 for p in dtensor.placements: 

76 if isinstance(p, Shard): 

77 shard_dim = p.dim if p.dim >= 0 else len(dtensor.shape) + p.dim 

78 break 

79 return mesh, shard_dim 

80 

81 

82def _get_local_tensor(tensor): 

83 """Get DTensor local tensor, or return regular tensor.""" 

84 if _is_dtensor(tensor): 

85 return tensor._local_tensor # type: ignore # pylint: disable=W0212 

86 return tensor 

87 

88 

89def _get_full_tensor(tensor): 

90 """Get full tensor (all_gather).""" 

91 if _is_dtensor(tensor): 

92 return tensor.full_tensor() # type: ignore 

93 return tensor 

94 

95 

96def _validate_target_type_base(is_floating: bool) -> None: 

97 """Validate target type (base check). 

98 

99 Args: 

100 is_floating: Whether target tensor has floating dtype. 

101 

102 Raises: 

103 ValueError: If target is floating point (probabilistic). 

104 """ 

105 if is_floating: 

106 raise ValueError( 

107 "Probabilistic target (float) is not supported in loss_parallel. " 

108 "Target must be class indices (int64)." 

109 ) 

110 

111 

112def _validate_mesh_and_shard(dtensor: DTensor, strict: bool = True) -> None: 

113 """Validate mesh and Shard dimension. 

114 

115 Args: 

116 dtensor: Input DTensor. 

117 strict: Whether to raise error or warn on validation failure. 

118 

119 Raises: 

120 ValueError: If mesh is not 1D or input is not Shard(-1) in strict mode. 

121 """ 

122 mesh = dtensor.device_mesh 

123 

124 if mesh.ndim != 1: 

125 if strict: 

126 raise ValueError( 

127 f"Expected 1D TP mesh, got {mesh.ndim}D mesh. " 

128 "Slice a 1D sub-mesh first: mesh['tp']" 

129 ) 

130 warnings.warn( 

131 f"Expected 1D TP mesh, got {mesh.ndim}D mesh. " 

132 "This may cause incorrect results." 

133 ) 

134 

135 if not _is_shard_on_last_dim(dtensor): 

136 if strict: 

137 raise ValueError( 

138 "Expected Shard(-1) on class dimension. " 

139 f"Got placements: {dtensor.placements}" 

140 ) 

141 warnings.warn( 

142 f"Expected Shard(-1) on class dimension. " 

143 f"Got placements: {dtensor.placements}. " 

144 "This may cause incorrect results." 

145 ) 

146 

147 

148def _validate_cross_entropy_params( 

149 input_tensor, 

150 target, 

151 weight: Optional[Any], 

152 size_average: Optional[bool], 

153 ignore_index: int, # pylint: disable=W0613 

154 reduce: Optional[bool], 

155 reduction: str, 

156 label_smoothing: float, 

157 is_floating_fn, 

158) -> None: 

159 """Validate cross_entropy parameter support. 

160 

161 Args: 

162 input_tensor: Input tensor. 

163 target: Target tensor. 

164 weight: Optional weight tensor. 

165 size_average: Deprecated parameter. 

166 ignore_index: Index to ignore (reserved for future use). 

167 reduce: Deprecated parameter. 

168 reduction: Reduction method. 

169 label_smoothing: Label smoothing factor. 

170 is_floating_fn: Platform-specific function to check if tensor is floating. 

171 

172 Raises: 

173 ValueError: If parameters are invalid. 

174 """ 

175 _validate_target_type_base(is_floating_fn(target)) 

176 

177 if _is_dtensor(input_tensor): 

178 if not _is_shard_on_last_dim(input_tensor): 

179 raise ValueError( 

180 "input must be Shard(-1) on class dimension. " 

181 f"Got placements: {input_tensor.placements}" 

182 ) 

183 else: 

184 raise ValueError( 

185 "input must be a DTensor when using loss_parallel. " 

186 f"Got type: {type(input)}" 

187 ) 

188 

189 if weight is not None and _is_dtensor(weight): 

190 if not _is_replicate(weight): 

191 raise ValueError( 

192 "weight must be Replicate when it's a DTensor. " 

193 f"Got placements: {weight.placements}" 

194 ) 

195 

196 if size_average is not None or reduce is not None: 

197 warnings.warn( 

198 "size_average and reduce arguments are deprecated. " 

199 "Please use reduction='mean' or reduction='sum' instead.", 

200 DeprecationWarning, 

201 ) 

202 

203 if label_smoothing != 0.0: 

204 raise ValueError( 

205 "label_smoothing is not supported in loss_parallel. " 

206 "Please set label_smoothing=0.0 or disable loss_parallel." 

207 ) 

208 

209 if reduction not in ("none", "mean", "sum"): 

210 raise ValueError(f"Invalid reduction: {reduction}. Must be 'none', 'mean', or 'sum'.") 

211 

212 

213def _check_context_and_layout(dtensor: DTensor) -> None: # pylint: disable=W0613 

214 """Check context and layout. 

215 

216 Args: 

217 dtensor: Input DTensor (reserved for future layout validation). 

218 

219 Raises: 

220 ValueError: If not in loss_parallel context. 

221 """ 

222 if not is_loss_parallel_active(): 

223 raise ValueError( 

224 "Shard logits detected but not in loss_parallel context. " 

225 "Please wrap with loss_parallel() context manager." 

226 )