Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / mindspore / loss_parallel_ops.py: 0%
190 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
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"""MindSpore-specific loss_parallel operations.
17Distributed cross-entropy kernel implementation using mindspore.common._grad_function._Function.
18"""
20from __future__ import annotations
22from typing import Any, Optional, Tuple
24import mindspore as ms # pylint: disable=C0415
25from mindspore import mint # pylint: disable=C0415
26from mindspore.common._grad_function import _Function # pylint: disable=C0415
27from mindspore.common.tensor import Tensor # pylint: disable=C0415
29from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
30from hyper_parallel.core.tensor_parallel.loss_parallel_ops_common import (
31 _is_dtensor,
32 _is_shard_on_last_dim,
33 _get_mesh_and_dim,
34 _get_local_tensor,
35 _validate_cross_entropy_params,
36 _check_context_and_layout,
37 _validate_mesh_and_shard,
38)
39from hyper_parallel.core.tensor_parallel.loss_parallel import _get_loss_parallel_strict
40from hyper_parallel.platform import get_platform
42platform = get_platform()
44__all__ = [
45 "distributed_cross_entropy",
46 "distributed_log_softmax",
47 "distributed_nll_loss_forward",
48 "DistributedCrossEntropyFunction",
49]
52def _is_floating_ms(tensor: Tensor) -> bool:
53 """Check if MindSpore tensor is floating point."""
54 return tensor.dtype in (ms.float16, ms.float32, ms.float64)
57def _compute_vocab_start(vocab_size: int, tp_size: int, rank: int) -> int:
58 """Compute the starting index for this rank's vocab shard.
60 Args:
61 vocab_size: Total vocabulary size.
62 tp_size: Tensor parallel world size.
63 rank: Current rank in TP mesh.
65 Returns:
66 Starting index of this rank's vocab shard.
68 Note:
69 This follows torch.chunk behavior: chunk_size = ceil(vocab_size/tp_size),
70 and each rank's start = rank * chunk_size. The last rank may have fewer elements.
71 """
72 chunk_size = (vocab_size + tp_size - 1) // tp_size # ceil division
73 return rank * chunk_size
76def distributed_log_softmax(
77 logits_local: Tensor,
78 dim: int,
79 mesh: DeviceMesh,
80 mesh_dim: int = 0,
81) -> Tensor:
82 """K1: Stable log-softmax on class-sharded dimension.
84 Args:
85 logits_local: Local logits shard.
86 dim: Class dimension.
87 mesh: DeviceMesh.
88 mesh_dim: Mesh dimension (default 0).
90 Returns:
91 Local log-softmax with unchanged layout.
93 Communication:
94 MAX + SUM all_reduce
95 """
96 max_local = logits_local.max(axis=dim, keepdims=True)
98 group = mesh.get_group(mesh_dim)
99 max_global = platform.differentiable_all_reduce(max_local, op="max", group=group)
101 exp_local = (logits_local - max_global).exp()
103 sum_local = exp_local.sum(axis=dim, keepdims=True)
105 sum_global = platform.differentiable_all_reduce(sum_local, op="sum", group=group)
107 log_softmax = logits_local - max_global - sum_global.log()
109 return log_softmax
112def distributed_nll_loss_forward(
113 log_probs: Tensor,
114 target: Tensor,
115 weight: Optional[Tensor],
116 ignore_index: int,
117 reduction: str,
118 vocab_start: int,
119 vocab_end: int,
120) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
121 """K2: Index target + optional weight + reduction.
123 Args:
124 log_probs: Sharded log_probs.
125 target: Target class indices.
126 weight: Optional weights.
127 ignore_index: Index to ignore.
128 reduction: Reduction method.
129 vocab_start: Start index of this vocab shard.
130 vocab_end: End index of this vocab shard.
132 Returns:
133 Tuple of (loss, total_weight, target_mask, vocab_start_tensor).
134 """
135 batch_size = target.numel()
137 target_flat = target.flatten()
139 target_mask = (target_flat >= vocab_start) & (target_flat < vocab_end)
141 ignore_mask = target_flat != ignore_index
142 target_mask = target_mask & ignore_mask
144 if reduction == "none":
145 loss = mint.zeros((batch_size,), dtype=log_probs.dtype)
146 else:
147 loss = mint.zeros((1,), dtype=log_probs.dtype)
149 total_weight = mint.zeros((1,), dtype=log_probs.dtype)
151 if target_mask.any():
152 local_target = target_flat[target_mask] - vocab_start
154 log_probs_2d = log_probs.reshape(-1, log_probs.shape[-1])
156 row_indices = mint.nonzero(target_mask).flatten()
158 selected_log_probs = log_probs_2d[row_indices, local_target]
160 if weight is not None:
161 global_target = target_flat[target_mask]
162 sample_weights = weight[global_target]
163 selected_log_probs = selected_log_probs * sample_weights
164 total_weight = sample_weights.sum().reshape(1)
165 else:
166 total_weight = ms.Tensor(
167 target_mask.sum().asnumpy().item(), dtype=log_probs.dtype
168 ).reshape(1)
170 nll = -selected_log_probs
172 if reduction == "none":
173 loss_flat = mint.zeros((batch_size,), dtype=log_probs.dtype)
174 loss_flat[target_mask] = nll
175 loss = loss_flat.reshape(target.shape)
176 elif reduction == "sum":
177 loss = nll.sum().unsqueeze(0)
178 else:
179 loss = nll.sum().unsqueeze(0)
180 else:
181 if reduction == "none":
182 loss = mint.zeros((batch_size,), dtype=log_probs.dtype).reshape(target.shape)
183 total_weight = mint.zeros((1,), dtype=log_probs.dtype)
185 vocab_start_tensor = ms.Tensor(vocab_start, dtype=ms.int64)
186 return loss, total_weight, target_mask, vocab_start_tensor
189class DistributedCrossEntropyFunction(_Function):
190 """K3: Fused backward for distributed cross_entropy (MindSpore version)."""
192 @staticmethod
193 def forward(
194 ctx: Any,
195 input_local: Tensor,
196 target: Tensor,
197 weight: Optional[Tensor],
198 ignore_index: int,
199 reduction: str,
200 vocab_size: int,
201 mesh: DeviceMesh,
202 mesh_dim: int,
203 ) -> Tensor:
204 """Forward pass."""
205 local_vocab_size = input_local.shape[-1]
206 rank = mesh.get_local_rank(mesh_dim)
207 tp_size = mesh.size(mesh_dim)
208 vocab_start = _compute_vocab_start(vocab_size, tp_size, rank)
209 vocab_end = vocab_start + local_vocab_size
211 log_probs_local = distributed_log_softmax(
212 input_local, dim=-1, mesh=mesh, mesh_dim=mesh_dim
213 )
215 loss, total_weight, target_mask, vocab_start_tensor = distributed_nll_loss_forward(
216 log_probs_local,
217 target,
218 weight,
219 ignore_index,
220 reduction,
221 vocab_start,
222 vocab_end,
223 )
225 if reduction == "mean":
226 group = mesh.get_group(mesh_dim)
227 total_loss = platform.differentiable_all_reduce(loss, op="sum", group=group)
228 total_weight_sum = platform.differentiable_all_reduce(
229 total_weight, op="sum", group=group
230 )
232 ctx.save_for_backward(
233 input_local,
234 log_probs_local,
235 target,
236 weight,
237 total_weight_sum,
238 target_mask,
239 vocab_start_tensor,
240 )
241 ctx.reduction = reduction
242 ctx.ignore_index = ignore_index
243 ctx.vocab_size = vocab_size
244 ctx.local_vocab_size = local_vocab_size
245 ctx.mesh = mesh
246 ctx.mesh_dim = mesh_dim
247 ctx.vocab_start = vocab_start
248 ctx.vocab_end = vocab_end
250 return total_loss / total_weight_sum.clamp(min=1e-12)
251 if reduction == "sum":
252 group = mesh.get_group(mesh_dim)
253 total_loss = platform.differentiable_all_reduce(loss, op="sum", group=group)
255 ctx.save_for_backward(
256 input_local,
257 log_probs_local,
258 target,
259 weight,
260 mint.zeros((1,), dtype=loss.dtype),
261 target_mask,
262 vocab_start_tensor,
263 )
264 ctx.reduction = reduction
265 ctx.ignore_index = ignore_index
266 ctx.vocab_size = vocab_size
267 ctx.local_vocab_size = local_vocab_size
268 ctx.mesh = mesh
269 ctx.mesh_dim = mesh_dim
270 ctx.vocab_start = vocab_start
271 ctx.vocab_end = vocab_end
273 return total_loss
274 ctx.save_for_backward(
275 input_local,
276 log_probs_local,
277 target,
278 weight,
279 mint.zeros((1,), dtype=loss.dtype),
280 target_mask,
281 vocab_start_tensor,
282 )
283 ctx.reduction = reduction
284 ctx.ignore_index = ignore_index
285 ctx.vocab_size = vocab_size
286 ctx.local_vocab_size = local_vocab_size
287 ctx.mesh = mesh
288 ctx.mesh_dim = mesh_dim
289 ctx.vocab_start = vocab_start
290 ctx.vocab_end = vocab_end
292 return loss
294 @staticmethod
295 def backward(ctx: Any, grad_output: Tensor) -> Tuple[Optional[Tensor], ...]:
296 """Backward pass (vectorized implementation)."""
297 (
298 _,
299 log_probs_local,
300 target,
301 weight,
302 total_weight,
303 _,
304 _,
305 ) = ctx.saved_tensors
307 reduction = ctx.reduction
308 ignore_index = ctx.ignore_index
309 _ = ctx.local_vocab_size
310 vocab_start = ctx.vocab_start
311 vocab_end = ctx.vocab_end
312 _ = ctx.mesh
313 _ = ctx.mesh_dim
315 batch_size = target.numel()
316 target_flat = target.flatten()
318 softmax_local = log_probs_local.exp()
320 ignore_mask = target_flat != ignore_index
322 if weight is not None:
323 sample_weights = weight[target_flat]
324 else:
325 sample_weights = None
327 if reduction == "mean":
328 grad_scale = grad_output / total_weight.clamp(min=1e-12)
329 elif reduction == "sum":
330 grad_scale = grad_output
331 else:
332 grad_scale = grad_output.flatten()
334 in_vocab_mask = (target_flat >= vocab_start) & (target_flat < vocab_end) & ignore_mask
336 if reduction == "none":
337 grad_scale_expanded = grad_scale.unsqueeze(-1)
338 if sample_weights is not None:
339 grad_scale_expanded = grad_scale_expanded * sample_weights.unsqueeze(-1)
340 grad_input = softmax_local * grad_scale_expanded
341 else:
342 if sample_weights is not None:
343 grad_scale = grad_scale * sample_weights.unsqueeze(-1)
344 grad_input = softmax_local * grad_scale.unsqueeze(-1)
346 local_targets = mint.where(in_vocab_mask, target_flat - vocab_start, mint.zeros_like(target_flat))
348 if in_vocab_mask.any():
349 row_indices = mint.arange(batch_size, dtype=ms.int64)
351 if reduction == "none":
352 if sample_weights is not None:
353 grad_values = -grad_scale * sample_weights
354 else:
355 grad_values = -grad_scale
356 else:
357 grad_values = -grad_scale.expand_as(target_flat)
359 grad_input = grad_input.contiguous()
360 grad_input[row_indices[in_vocab_mask], local_targets[in_vocab_mask]] += grad_values[in_vocab_mask]
362 if not ignore_mask.all():
363 ignore_indices = ~ignore_mask
364 if reduction == "none":
365 grad_input[ignore_indices] = 0.0
366 else:
367 ignore_indices_expanded = ignore_indices.unsqueeze(-1).expand_as(grad_input)
368 grad_input = mint.where(ignore_indices_expanded, grad_input, mint.zeros_like(grad_input))
370 return grad_input, None, None, None, None, None, None, None
373def distributed_cross_entropy(
374 input_tensor: Tensor,
375 target: Tensor,
376 weight: Optional[Tensor] = None,
377 size_average: Optional[bool] = None,
378 ignore_index: int = -100,
379 reduce: Optional[bool] = None,
380 reduction: str = "mean",
381 label_smoothing: float = 0.0,
382) -> Tensor:
383 """Distributed cross_entropy main entry (MindSpore version).
385 Args:
386 input_tensor: Must be DTensor with Shard(-1) on class dimension.
387 target: Class indices with Replicate or consistent layout.
388 weight: Optional weights; if DTensor, must be Replicate.
389 size_average: Deprecated.
390 ignore_index: Index to ignore (default -100).
391 reduce: Deprecated.
392 reduction: Reduction method: 'none', 'mean', or 'sum'.
393 label_smoothing: Not supported (must be 0.0).
395 Returns:
396 Loss tensor.
397 """
398 input_dtensor = None
399 mesh = None
400 vocab_size = None
402 if _is_dtensor(input_tensor):
403 if not _is_shard_on_last_dim(input_tensor):
404 raise ValueError(
405 "input must be Shard(-1) on class dimension. "
406 f"Got placements: {input_tensor.placements}"
407 )
408 input_dtensor = input_tensor
409 mesh, _ = _get_mesh_and_dim(input_tensor)
410 vocab_size = input_tensor.shape[-1]
412 input_for_check = input_dtensor if input_dtensor is not None else input_tensor
413 _check_context_and_layout(input_for_check) # type: ignore
415 _validate_cross_entropy_params(
416 input_tensor,
417 target,
418 weight,
419 size_average,
420 ignore_index,
421 reduce,
422 reduction,
423 label_smoothing,
424 _is_floating_ms,
425 )
427 if input_dtensor is not None:
428 input_local = _get_local_tensor(input_dtensor)
429 local_vocab_size = input_local.shape[-1]
431 if input_dtensor.ndim > 2:
432 input_local = input_local.reshape(-1, local_vocab_size)
433 target = target.reshape(-1)
434 else:
435 raise ValueError(
436 "input must be a DTensor when using loss_parallel. "
437 f"Got type: {type(input_tensor)}"
438 )
440 strict = _get_loss_parallel_strict()
441 _validate_mesh_and_shard(input_dtensor, strict) # type: ignore
443 mesh_dim = 0
445 loss = DistributedCrossEntropyFunction.apply(
446 input_local,
447 target,
448 weight,
449 ignore_index,
450 reduction,
451 vocab_size,
452 mesh,
453 mesh_dim,
454 )
456 return loss
459def distributed_cross_entropy_from_op_call(
460 op_call: Any, # pylint: disable=W0613
461 args: tuple,
462 kwargs: dict,
463) -> Tensor:
464 """Parse arguments from op_call and invoke distributed cross_entropy.
466 Used for OpDispatcher routing.
468 Args:
469 op_call: Op call object (reserved for future use).
470 args: Positional arguments.
471 kwargs: Keyword arguments.
473 Returns:
474 Loss tensor.
475 """
476 input_tensor = args[0] if len(args) > 0 else kwargs.get("input")
477 target = args[1] if len(args) > 1 else kwargs.get("target")
478 weight = args[2] if len(args) > 2 else kwargs.get("weight")
479 size_average = args[3] if len(args) > 3 else kwargs.get("size_average")
480 ignore_index = args[4] if len(args) > 4 else kwargs.get("ignore_index", -100)
481 reduce = args[5] if len(args) > 5 else kwargs.get("reduce")
482 reduction = args[6] if len(args) > 6 else kwargs.get("reduction", "mean")
483 label_smoothing = args[7] if len(args) > 7 else kwargs.get("label_smoothing", 0.0)
485 return distributed_cross_entropy(
486 input_tensor=input_tensor,
487 target=target,
488 weight=weight,
489 size_average=size_average,
490 ignore_index=ignore_index,
491 reduce=reduce,
492 reduction=reduction,
493 label_smoothing=label_smoothing,
494 )