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 platform-agnostic helper functions and validation logic 

18shared between PyTorch and MindSpore implementations. 

19""" 

20 

21from __future__ import annotations 

22 

23import warnings 

24from typing import Any, Optional, TYPE_CHECKING 

25 

26from hyper_parallel.core.dtensor.dtensor import DTensor 

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

28from hyper_parallel.core.tensor_parallel.loss_parallel import is_loss_parallel_active 

29 

30if TYPE_CHECKING: 

31 from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

32 

33__all__ = [ 

34 "_is_dtensor", 

35 "_is_shard_on_last_dim", 

36 "_is_replicate", 

37 "_get_mesh_and_dim", 

38 "_get_local_tensor", 

39 "_get_full_tensor", 

40 "_validate_target_type_base", 

41 "_validate_mesh_and_shard", 

42 "_validate_cross_entropy_params", 

43 "_check_context_and_layout", 

44] 

45 

46 

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

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

49 return isinstance(obj, DTensor) 

50 

51 

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

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

54 if not _is_dtensor(dtensor): 

55 return False 

56 placements = dtensor.placements 

57 if not placements: 

58 return False 

59 last_placement = placements[-1] 

60 if isinstance(last_placement, Shard): 

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

62 return False 

63 

64 

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

66 """Check if DTensor is Replicate.""" 

67 if not _is_dtensor(dtensor): 

68 return False 

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

70 

71 

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

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

74 mesh = dtensor.device_mesh 

75 shard_dim = None 

76 for p in dtensor.placements: 

77 if isinstance(p, Shard): 

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

79 break 

80 return mesh, shard_dim 

81 

82 

83def _get_local_tensor(tensor): 

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

85 if _is_dtensor(tensor): 

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

87 return tensor 

88 

89 

90def _get_full_tensor(tensor): 

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

92 if _is_dtensor(tensor): 

93 return tensor.full_tensor() # type: ignore 

94 return tensor 

95 

96 

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

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

99 

100 Args: 

101 is_floating: Whether target tensor has floating dtype. 

102 

103 Raises: 

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

105 """ 

106 if is_floating: 

107 raise ValueError( 

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

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

110 ) 

111 

112 

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

114 """Validate mesh and Shard dimension. 

115 

116 Args: 

117 dtensor: Input DTensor. 

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

119 

120 Raises: 

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

122 """ 

123 mesh = dtensor.device_mesh 

124 

125 if mesh.ndim != 1: 

126 if strict: 

127 raise ValueError( 

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

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

130 ) 

131 warnings.warn( 

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

133 "This may cause incorrect results." 

134 ) 

135 

136 if not _is_shard_on_last_dim(dtensor): 

137 if strict: 

138 raise ValueError( 

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

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

141 ) 

142 warnings.warn( 

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

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

145 "This may cause incorrect results." 

146 ) 

147 

148 

149def _validate_cross_entropy_params( 

150 input_tensor, 

151 target, 

152 weight: Optional[Any], 

153 size_average: Optional[bool], 

154 ignore_index: int, # pylint: disable=W0613 

155 reduce: Optional[bool], 

156 reduction: str, 

157 label_smoothing: float, 

158 is_floating_fn, 

159) -> None: 

160 """Validate cross_entropy parameter support. 

161 

162 Args: 

163 input_tensor: Input tensor. 

164 target: Target tensor. 

165 weight: Optional weight tensor. 

166 size_average: Deprecated parameter. 

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

168 reduce: Deprecated parameter. 

169 reduction: Reduction method. 

170 label_smoothing: Label smoothing factor. 

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

172 

173 Raises: 

174 ValueError: If parameters are invalid. 

175 """ 

176 _validate_target_type_base(is_floating_fn(target)) 

177 

178 if _is_dtensor(input_tensor): 

179 if not _is_shard_on_last_dim(input_tensor): 

180 raise ValueError( 

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

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

183 ) 

184 else: 

185 raise ValueError( 

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

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

188 ) 

189 

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

191 if not _is_replicate(weight): 

192 raise ValueError( 

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

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

195 ) 

196 

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

198 warnings.warn( 

199 "size_average and reduce arguments are deprecated. " 

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

201 DeprecationWarning, 

202 ) 

203 

204 if label_smoothing != 0.0: 

205 raise ValueError( 

206 "label_smoothing is not supported in loss_parallel. " 

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

208 ) 

209 

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

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

212 

213 

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

215 """Check context and layout. 

216 

217 Args: 

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

219 

220 Raises: 

221 ValueError: If not in loss_parallel context. 

222 """ 

223 if not is_loss_parallel_active(): 

224 raise ValueError( 

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

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

227 )