Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / loss_parallel_ops.py: 0%
188 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"""PyTorch-specific loss_parallel operations.
17Distributed cross-entropy kernel implementation using torch.autograd.Function.
18"""
20from __future__ import annotations
22from typing import Any, Optional, Tuple
24import torch # pylint: disable=C0415
25from torch import Tensor # pylint: disable=C0415
27from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
28from hyper_parallel.core.tensor_parallel.loss_parallel_ops_common import (
29 _is_dtensor,
30 _is_shard_on_last_dim,
31 _get_mesh_and_dim,
32 _get_local_tensor,
33 _validate_cross_entropy_params,
34 _check_context_and_layout,
35 _validate_mesh_and_shard,
36)
37from hyper_parallel.core.tensor_parallel.loss_parallel import _get_loss_parallel_strict
38from hyper_parallel.platform import get_platform
40platform = get_platform()
42__all__ = [
43 "distributed_cross_entropy",
44 "distributed_log_softmax",
45 "distributed_nll_loss_forward",
46 "DistributedCrossEntropyFunction",
47]
50def _is_floating_torch(tensor: Tensor) -> bool:
51 """Check if PyTorch tensor is floating point."""
52 return tensor.is_floating_point()
55def _compute_vocab_start(vocab_size: int, tp_size: int, rank: int) -> int:
56 """Compute the starting index for this rank's vocab shard.
58 Args:
59 vocab_size: Total vocabulary size.
60 tp_size: Tensor parallel world size.
61 rank: Current rank in TP mesh.
63 Returns:
64 Starting index of this rank's vocab shard.
66 Note:
67 This follows torch.chunk behavior: chunk_size = ceil(vocab_size/tp_size),
68 and each rank's start = rank * chunk_size. The last rank may have fewer elements.
69 """
70 chunk_size = (vocab_size + tp_size - 1) // tp_size # ceil division
71 return rank * chunk_size
74def distributed_log_softmax(
75 logits_local: Tensor,
76 dim: int,
77 mesh: DeviceMesh,
78 mesh_dim: int = 0,
79) -> Tensor:
80 """K1: Stable log-softmax on class-sharded dimension.
82 Args:
83 logits_local: Local logits shard.
84 dim: Class dimension.
85 mesh: DeviceMesh.
86 mesh_dim: Mesh dimension (default 0).
88 Returns:
89 Local log-softmax with unchanged layout.
91 Communication:
92 MAX + SUM all_reduce
93 """
94 max_local = logits_local.max(dim=dim, keepdim=True).values
96 group = mesh.get_group(mesh_dim)
97 max_global = platform.differentiable_all_reduce(max_local, op="max", group=group)
99 exp_local = (logits_local - max_global).exp()
101 sum_local = exp_local.sum(dim=dim, keepdim=True)
103 sum_global = platform.differentiable_all_reduce(sum_local, op="sum", group=group)
105 log_softmax = logits_local - max_global - sum_global.log()
107 return log_softmax
110def distributed_nll_loss_forward(
111 log_probs: Tensor,
112 target: Tensor,
113 weight: Optional[Tensor],
114 ignore_index: int,
115 reduction: str,
116 vocab_start: int,
117 vocab_end: int,
118) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
119 """K2: Index target + optional weight + reduction.
121 Args:
122 log_probs: Sharded log_probs.
123 target: Target class indices.
124 weight: Optional weights.
125 ignore_index: Index to ignore.
126 reduction: Reduction method.
127 vocab_start: Start index of this vocab shard.
128 vocab_end: End index of this vocab shard.
130 Returns:
131 Tuple of (loss, total_weight, target_mask, vocab_start_tensor).
132 """
133 batch_size = target.numel()
135 target_flat = target.flatten()
137 target_mask = (target_flat >= vocab_start) & (target_flat < vocab_end)
139 ignore_mask = target_flat != ignore_index
140 target_mask = target_mask & ignore_mask
142 if reduction == "none":
143 loss = torch.zeros(batch_size, dtype=log_probs.dtype, device=log_probs.device)
144 else:
145 loss = torch.zeros(1, dtype=log_probs.dtype, device=log_probs.device)
147 total_weight = torch.zeros(1, dtype=log_probs.dtype, device=log_probs.device)
149 if target_mask.any():
150 local_target = target_flat[target_mask] - vocab_start
152 log_probs_2d = log_probs.reshape(-1, log_probs.shape[-1])
154 row_indices = torch.where(target_mask)[0]
156 selected_log_probs = log_probs_2d[row_indices, local_target]
158 if weight is not None:
159 global_target = target_flat[target_mask]
160 sample_weights = weight[global_target]
161 selected_log_probs = selected_log_probs * sample_weights
162 total_weight = sample_weights.sum().reshape(1)
163 else:
164 total_weight = torch.tensor(
165 target_mask.sum().item(), dtype=log_probs.dtype, device=log_probs.device
166 ).reshape(1)
168 nll = -selected_log_probs
170 if reduction == "none":
171 loss_flat = torch.zeros(batch_size, dtype=log_probs.dtype, device=log_probs.device)
172 loss_flat[target_mask] = nll
173 loss = loss_flat.reshape(target.shape)
174 elif reduction == "sum":
175 loss = nll.sum().unsqueeze(0)
176 else:
177 loss = nll.sum().unsqueeze(0)
178 else:
179 if reduction == "none":
180 loss = torch.zeros(
181 batch_size, dtype=log_probs.dtype, device=log_probs.device
182 ).reshape(target.shape)
183 total_weight = torch.zeros(1, dtype=log_probs.dtype, device=log_probs.device)
185 return loss, total_weight, target_mask, torch.tensor(
186 vocab_start, dtype=torch.long, device=log_probs.device
187 )
190class DistributedCrossEntropyFunction(torch.autograd.Function):
191 """K3: Fused backward for distributed cross_entropy."""
193 @staticmethod
194 def forward(
195 ctx: Any,
196 input_local: Tensor,
197 target: Tensor,
198 weight: Optional[Tensor],
199 ignore_index: int,
200 reduction: str,
201 vocab_size: int,
202 mesh: DeviceMesh,
203 mesh_dim: int,
204 ) -> Tensor:
205 """Forward pass."""
206 local_vocab_size = input_local.shape[-1]
207 rank = mesh.get_local_rank(mesh_dim)
208 tp_size = mesh.size(mesh_dim)
209 vocab_start = _compute_vocab_start(vocab_size, tp_size, rank)
210 vocab_end = vocab_start + local_vocab_size
212 log_probs_local = distributed_log_softmax(
213 input_local, dim=-1, mesh=mesh, mesh_dim=mesh_dim
214 )
216 loss, total_weight, target_mask, vocab_start_tensor = distributed_nll_loss_forward(
217 log_probs_local,
218 target,
219 weight,
220 ignore_index,
221 reduction,
222 vocab_start,
223 vocab_end,
224 )
226 if reduction == "mean":
227 group = mesh.get_group(mesh_dim)
228 total_loss = platform.differentiable_all_reduce(loss, op="sum", group=group)
229 total_weight_sum = platform.differentiable_all_reduce(
230 total_weight, op="sum", group=group
231 )
233 ctx.save_for_backward(
234 input_local,
235 log_probs_local,
236 target,
237 weight,
238 total_weight_sum,
239 target_mask,
240 vocab_start_tensor,
241 )
242 ctx.reduction = reduction
243 ctx.ignore_index = ignore_index
244 ctx.vocab_size = vocab_size
245 ctx.local_vocab_size = local_vocab_size
246 ctx.mesh = mesh
247 ctx.mesh_dim = mesh_dim
248 ctx.vocab_start = vocab_start
249 ctx.vocab_end = vocab_end
251 if total_weight_sum.item() == 0:
252 return torch.tensor(float('nan'), dtype=total_loss.dtype, device=total_loss.device)
253 return total_loss / total_weight_sum
254 if reduction == "sum":
255 group = mesh.get_group(mesh_dim)
256 total_loss = platform.differentiable_all_reduce(loss, op="sum", group=group)
258 ctx.save_for_backward(
259 input_local,
260 log_probs_local,
261 target,
262 weight,
263 torch.zeros(1, dtype=loss.dtype, device=loss.device),
264 target_mask,
265 vocab_start_tensor,
266 )
267 ctx.reduction = reduction
268 ctx.ignore_index = ignore_index
269 ctx.vocab_size = vocab_size
270 ctx.local_vocab_size = local_vocab_size
271 ctx.mesh = mesh
272 ctx.mesh_dim = mesh_dim
273 ctx.vocab_start = vocab_start
274 ctx.vocab_end = vocab_end
276 return total_loss
277 ctx.save_for_backward(
278 input_local,
279 log_probs_local,
280 target,
281 weight,
282 torch.zeros(1, dtype=loss.dtype, device=loss.device),
283 target_mask,
284 vocab_start_tensor,
285 )
286 ctx.reduction = reduction
287 ctx.ignore_index = ignore_index
288 ctx.vocab_size = vocab_size
289 ctx.local_vocab_size = local_vocab_size
290 ctx.mesh = mesh
291 ctx.mesh_dim = mesh_dim
292 ctx.vocab_start = vocab_start
293 ctx.vocab_end = vocab_end
295 return loss
297 @staticmethod
298 def backward(ctx: Any, grad_output: Tensor) -> Tuple[Optional[Tensor], ...]:
299 """Backward pass (vectorized implementation)."""
300 (
301 _,
302 log_probs_local,
303 target,
304 weight,
305 total_weight,
306 _,
307 _,
308 ) = ctx.saved_tensors
310 reduction = ctx.reduction
311 ignore_index = ctx.ignore_index
312 _ = ctx.local_vocab_size
313 vocab_start = ctx.vocab_start
314 vocab_end = ctx.vocab_end
315 _ = ctx.mesh
316 _ = ctx.mesh_dim
318 batch_size = target.numel()
319 target_flat = target.flatten()
321 softmax_local = log_probs_local.exp()
323 ignore_mask = target_flat != ignore_index
325 if weight is not None:
326 sample_weights = weight[target_flat]
327 else:
328 sample_weights = None
330 if reduction == "mean":
331 grad_scale = grad_output / total_weight.clamp(min=1e-12)
332 elif reduction == "sum":
333 grad_scale = grad_output
334 else:
335 grad_scale = grad_output.flatten()
337 in_vocab_mask = (target_flat >= vocab_start) & (target_flat < vocab_end) & ignore_mask
339 if reduction == "none":
340 grad_scale_expanded = grad_scale.unsqueeze(-1)
341 if sample_weights is not None:
342 grad_scale_expanded = grad_scale_expanded * sample_weights.unsqueeze(-1)
343 grad_input = softmax_local * grad_scale_expanded
344 else:
345 if sample_weights is not None:
346 grad_scale = grad_scale * sample_weights.unsqueeze(-1)
347 grad_input = softmax_local * grad_scale.unsqueeze(-1)
349 local_targets = torch.where(in_vocab_mask, target_flat - vocab_start, torch.zeros_like(target_flat))
351 if in_vocab_mask.any():
352 row_indices = torch.arange(batch_size, device=target.device, dtype=torch.long)
354 if reduction == "none":
355 if sample_weights is not None:
356 grad_values = -grad_scale * sample_weights
357 else:
358 grad_values = -grad_scale
359 else:
360 grad_values = -grad_scale.expand_as(target_flat)
362 grad_input = grad_input.contiguous()
363 grad_input[row_indices[in_vocab_mask], local_targets[in_vocab_mask]] += grad_values[in_vocab_mask]
365 if not ignore_mask.all():
366 if reduction == "none":
367 grad_input[~ignore_mask] = 0.0
368 else:
369 ignore_indices_expanded = (~ignore_mask).unsqueeze(-1).expand_as(grad_input)
370 grad_input[ignore_indices_expanded] = 0.0
372 return grad_input, None, None, None, None, None, None, None
375def distributed_cross_entropy(
376 input_tensor: Tensor,
377 target: Tensor,
378 weight: Optional[Tensor] = None,
379 size_average: Optional[bool] = None,
380 ignore_index: int = -100,
381 reduce: Optional[bool] = None,
382 reduction: str = "mean",
383 label_smoothing: float = 0.0,
384) -> Tensor:
385 """Distributed cross_entropy main entry (PyTorch version).
387 Args:
388 input_tensor: Must be DTensor with Shard(-1) on class dimension.
389 target: Class indices with Replicate or consistent layout.
390 weight: Optional weights; if DTensor, must be Replicate.
391 size_average: Deprecated.
392 ignore_index: Index to ignore (default -100).
393 reduce: Deprecated.
394 reduction: Reduction method: 'none', 'mean', or 'sum'.
395 label_smoothing: Not supported (must be 0.0).
397 Returns:
398 Loss tensor.
399 """
400 input_dtensor = None
401 mesh = None
402 vocab_size = None
404 if _is_dtensor(input_tensor):
405 if not _is_shard_on_last_dim(input_tensor):
406 raise ValueError(
407 "input must be Shard(-1) on class dimension. "
408 f"Got placements: {input_tensor.placements}"
409 )
410 input_dtensor = input_tensor
411 mesh, _ = _get_mesh_and_dim(input_tensor)
412 vocab_size = input_tensor.shape[-1]
414 input_for_check = input_dtensor if input_dtensor is not None else input_tensor
415 _check_context_and_layout(input_for_check) # type: ignore
417 _validate_cross_entropy_params(
418 input_tensor,
419 target,
420 weight,
421 size_average,
422 ignore_index,
423 reduce,
424 reduction,
425 label_smoothing,
426 _is_floating_torch,
427 )
429 if input_dtensor is not None:
430 input_local = _get_local_tensor(input_dtensor)
431 local_vocab_size = input_local.shape[-1]
433 if input_dtensor.ndim > 2:
434 input_local = input_local.reshape(-1, local_vocab_size)
435 target = target.reshape(-1)
436 else:
437 raise ValueError(
438 "input must be a DTensor when using loss_parallel. "
439 f"Got type: {type(input_tensor)}"
440 )
442 strict = _get_loss_parallel_strict()
443 _validate_mesh_and_shard(input_dtensor, strict) # type: ignore
445 mesh_dim = 0
447 loss = DistributedCrossEntropyFunction.apply(
448 input_local,
449 target,
450 weight,
451 ignore_index,
452 reduction,
453 vocab_size,
454 mesh,
455 mesh_dim,
456 )
458 return loss
461def distributed_cross_entropy_from_op_call(
462 op_call: Any, # pylint: disable=W0613
463 args: tuple,
464 kwargs: dict,
465) -> Tensor:
466 """Parse arguments from op_call and invoke distributed cross_entropy.
468 Used for OpDispatcher routing.
470 Args:
471 op_call: Op call object (reserved for future use).
472 args: Positional arguments.
473 kwargs: Keyword arguments.
475 Returns:
476 Loss tensor.
477 """
478 input_tensor = args[0] if len(args) > 0 else kwargs.get("input")
479 target = args[1] if len(args) > 1 else kwargs.get("target")
480 weight = args[2] if len(args) > 2 else kwargs.get("weight")
481 size_average = args[3] if len(args) > 3 else kwargs.get("size_average")
482 ignore_index = args[4] if len(args) > 4 else kwargs.get("ignore_index", -100)
483 reduce = args[5] if len(args) > 5 else kwargs.get("reduce")
484 reduction = args[6] if len(args) > 6 else kwargs.get("reduction", "mean")
485 label_smoothing = args[7] if len(args) > 7 else kwargs.get("label_smoothing", 0.0)
487 return distributed_cross_entropy(
488 input_tensor=input_tensor,
489 target=target,
490 weight=weight,
491 size_average=size_average,
492 ignore_index=ignore_index,
493 reduce=reduce,
494 reduction=reduction,
495 label_smoothing=label_smoothing,
496 )