Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / losses / _vocab_parallel_cross_entropy.py: 90%
161 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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"""AutoModels-private vocab-parallel cross entropy (local-first).
17Merges the former ``components/loss/loss_parallel_ops.py`` kernel with the
18``loss_parallel_ops_common.py`` validation helpers into a single private
19production implementation. The leading underscore marks this module as
20internal: the generic DTensor loss API stays in ``core.tensor_parallel``
21(frozen), while AutoModels production loss uses this module.
23The public-in-package entry is :func:`vocab_parallel_cross_entropy_local`:
24it takes the *local* logits shard plus the global vocabulary size, the mesh
25and the mesh axis that shards the class dimension — no temporary DTensor is
26created per call. The 2D mesh-axis recognition and the "class dimension
27sharded by exactly one mesh axis" validation from the former AutoModels
28DTensor wrapper are absorbed by :func:`_resolve_class_mesh_dim`.
29"""
31from __future__ import annotations
33__all__ = [
34 "vocab_parallel_cross_entropy_local",
35 "distributed_log_softmax",
36 "distributed_nll_loss_forward",
37 "DistributedCrossEntropyFunction",
38]
40from typing import Any, NamedTuple, Optional, Tuple, TYPE_CHECKING
42import torch
43import torch.distributed.nn.functional as dist_func
44from torch import Tensor
46if TYPE_CHECKING:
47 from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
49_REDUCE_OP_MAP = {
50 "sum": torch.distributed.ReduceOp.SUM,
51 "prod": torch.distributed.ReduceOp.PRODUCT,
52 "max": torch.distributed.ReduceOp.MAX,
53 "min": torch.distributed.ReduceOp.MIN,
54 "all": torch.distributed.ReduceOp.MIN,
55 "avg": torch.distributed.ReduceOp.SUM,
56}
59def _differentiable_all_reduce(data: Tensor, op: str, group: Any) -> Tensor:
60 """Autograd-aware all-reduce: backward is an identity, the loss is scaled by world size."""
61 if not data.is_contiguous() or data.storage_offset() != 0:
62 data = data.contiguous()
63 reduce_op = _REDUCE_OP_MAP.get(op, torch.distributed.ReduceOp.SUM)
64 return dist_func.all_reduce(data, op=reduce_op, group=group)
67def _validate_target_type_base(is_floating: bool) -> None:
68 """Validate target type (base check).
70 Args:
71 is_floating: Whether target tensor has floating dtype.
73 Raises:
74 ValueError: If target is floating point (probabilistic).
75 """
76 if is_floating:
77 raise ValueError(
78 "Probabilistic target (float) is not supported in loss_parallel. "
79 "Target must be class indices (int64)."
80 )
83def _resolve_class_mesh_dim(mesh: "DeviceMesh", mesh_dim: Optional[int]) -> int:
84 """Identify the single mesh axis that shards the class dimension.
86 Absorbs the 2D mesh-axis recognition of the former DTensor-based
87 implementation: an explicit ``mesh_dim`` is validated against the mesh,
88 a 1D mesh resolves to axis 0, and a multi-dimensional mesh resolves the
89 axis named ``"tp"``. The class dimension must be sharded by exactly one
90 mesh axis.
92 Args:
93 mesh: Mesh holding the class-sharding axis.
94 mesh_dim: Explicitly declared class-sharding axis, or ``None`` to
95 resolve it from the mesh shape and dimension names.
97 Returns:
98 The resolved mesh axis index.
100 Raises:
101 ValueError: If no unique class-sharding axis can be identified.
102 """
103 if mesh_dim is not None:
104 if not isinstance(mesh_dim, int) or not 0 <= mesh_dim < mesh.ndim:
105 raise ValueError(
106 "Expected the class dimension to be sharded on exactly one "
107 f"mesh dimension. Got mesh_dim={mesh_dim} for a "
108 f"{mesh.ndim}D mesh."
109 )
110 return mesh_dim
111 if mesh.ndim == 1:
112 return 0
113 names = mesh.mesh_dim_names or ()
114 matches = [axis for axis, name in enumerate(names) if name == "tp"]
115 if len(matches) != 1:
116 raise ValueError(
117 "Expected the class dimension to be sharded on exactly one mesh "
118 f"dimension. Got {mesh.ndim}D mesh with dimension names {names}; "
119 "pass mesh_dim explicitly or name the sharding axis 'tp'."
120 )
121 return matches[0]
124def _is_floating_torch(tensor: Tensor) -> bool:
125 """Check if PyTorch tensor is floating point."""
126 return tensor.is_floating_point()
129def _compute_vocab_start(vocab_size: int, tp_size: int, rank: int) -> int:
130 """Compute the starting index for this rank's vocab shard.
132 Args:
133 vocab_size: Total vocabulary size.
134 tp_size: Tensor parallel world size.
135 rank: Current rank in TP mesh.
137 Returns:
138 Starting index of this rank's vocab shard.
140 Note:
141 This follows torch.chunk behavior: chunk_size = ceil(vocab_size/tp_size),
142 and each rank's start = rank * chunk_size. The last rank may have fewer elements.
143 """
144 chunk_size = (vocab_size + tp_size - 1) // tp_size # ceil division
145 return rank * chunk_size
148def distributed_log_softmax(
149 logits_local: Tensor,
150 dim: int,
151 mesh: DeviceMesh,
152 mesh_dim: int = 0,
153) -> Tensor:
154 """K1: Stable log-softmax on class-sharded dimension.
156 Args:
157 logits_local: Local logits shard.
158 dim: Class dimension.
159 mesh: DeviceMesh.
160 mesh_dim: Mesh dimension (default 0).
162 Returns:
163 Local log-softmax with unchanged layout.
165 Communication:
166 MAX + SUM all_reduce
167 """
168 max_local = logits_local.max(dim=dim, keepdim=True).values
170 group = mesh.get_group(mesh_dim)
171 max_global = _differentiable_all_reduce(max_local, op="max", group=group)
173 exp_local = (logits_local - max_global).exp()
175 sum_local = exp_local.sum(dim=dim, keepdim=True)
177 sum_global = _differentiable_all_reduce(sum_local, op="sum", group=group)
179 log_softmax = logits_local - max_global - sum_global.log()
181 return log_softmax
184def distributed_nll_loss_forward(
185 log_probs: Tensor,
186 target: Tensor,
187 weight: Optional[Tensor],
188 ignore_index: int,
189 reduction: str,
190 vocab_start: int,
191 vocab_end: int,
192) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
193 """K2: Index target + optional weight + reduction.
195 Args:
196 log_probs: Sharded log_probs.
197 target: Target class indices.
198 weight: Optional weights.
199 ignore_index: Index to ignore.
200 reduction: Reduction method.
201 vocab_start: Start index of this vocab shard.
202 vocab_end: End index of this vocab shard.
204 Returns:
205 Tuple of (loss, total_weight, target_mask, vocab_start_tensor).
206 """
207 batch_size = target.numel()
209 target_flat = target.flatten()
211 target_mask = (target_flat >= vocab_start) & (target_flat < vocab_end)
213 target_mask = target_mask & (target_flat != ignore_index)
215 if reduction == "none":
216 loss = torch.zeros(batch_size, dtype=log_probs.dtype, device=log_probs.device)
217 else:
218 loss = torch.zeros(1, dtype=log_probs.dtype, device=log_probs.device)
220 total_weight = torch.zeros(1, dtype=log_probs.dtype, device=log_probs.device)
222 if target_mask.any():
223 local_target = target_flat[target_mask] - vocab_start
225 selected_log_probs = log_probs.reshape(-1, log_probs.shape[-1])[
226 torch.where(target_mask)[0], local_target
227 ]
229 if weight is not None:
230 sample_weights = weight[target_flat[target_mask]]
231 selected_log_probs = selected_log_probs * sample_weights
232 total_weight = sample_weights.sum().reshape(1)
233 else:
234 total_weight = torch.tensor(
235 target_mask.sum().item(), dtype=log_probs.dtype, device=log_probs.device
236 ).reshape(1)
238 selected_log_probs = -selected_log_probs
240 if reduction == "none":
241 loss = torch.zeros(batch_size, dtype=log_probs.dtype, device=log_probs.device)
242 loss[target_mask] = selected_log_probs
243 loss = loss.reshape(target.shape)
244 elif reduction == "sum":
245 loss = selected_log_probs.sum().unsqueeze(0)
246 else:
247 loss = selected_log_probs.sum().unsqueeze(0)
248 else:
249 if reduction == "none":
250 loss = torch.zeros(
251 batch_size, dtype=log_probs.dtype, device=log_probs.device
252 ).reshape(target.shape)
253 total_weight = torch.zeros(1, dtype=log_probs.dtype, device=log_probs.device)
255 return loss, total_weight, target_mask, torch.tensor(
256 vocab_start, dtype=torch.long, device=log_probs.device
257 )
260class _CrossEntropyGradients(NamedTuple):
261 """Gradients aligned with the eight inputs of the autograd function."""
263 input_local: Tensor
264 target: None = None
265 weight: None = None
266 ignore_index: None = None
267 reduction: None = None
268 vocab_size: None = None
269 mesh: None = None
270 mesh_dim: None = None
273class _CrossEntropyContext(NamedTuple):
274 """State required by the custom cross-entropy backward pass."""
276 log_probs_local: Tensor
277 target: Tensor
278 weight: Optional[Tensor]
279 total_weight: Tensor
280 target_mask: Tensor
281 vocab_start_tensor: Tensor
282 reduction: str
283 ignore_index: int
284 vocab_start: int
285 vocab_end: int
288def _compute_cross_entropy_gradient(
289 state: _CrossEntropyContext,
290 grad_output: Tensor,
291) -> Tensor:
292 """Compute the local logits gradient for one reduction mode."""
293 target_flat = state.target.flatten()
294 softmax_local = state.log_probs_local.exp()
295 ignore_mask = target_flat != state.ignore_index
297 if state.weight is not None:
298 # Ignore positions can contain -100, so use a valid class index before
299 # masking their gradients out below.
300 safe_target = torch.where(ignore_mask, target_flat, torch.zeros_like(target_flat))
301 sample_weights = state.weight[safe_target]
302 else:
303 sample_weights = None
305 if state.reduction == "mean":
306 grad_scale = grad_output / state.total_weight.clamp(min=1e-12)
307 elif state.reduction == "sum":
308 grad_scale = grad_output
309 else:
310 grad_scale = grad_output.flatten()
312 in_vocab_mask = (target_flat >= state.vocab_start) & (target_flat < state.vocab_end) & ignore_mask
313 grad_scale_expanded = (
314 grad_scale.unsqueeze(-1) if state.reduction == "none" else grad_scale.reshape(1, 1)
315 )
316 if sample_weights is not None:
317 grad_scale_expanded = grad_scale_expanded * sample_weights.unsqueeze(-1)
318 grad_input = softmax_local * grad_scale_expanded
319 local_targets = torch.where(
320 in_vocab_mask, target_flat - state.vocab_start, torch.zeros_like(target_flat)
321 )
323 if in_vocab_mask.any():
324 row_indices = torch.arange(target_flat.numel(), device=state.target.device, dtype=torch.long)
325 if state.reduction == "none":
326 grad_values = -grad_scale
327 else:
328 grad_values = -grad_scale.expand_as(target_flat)
329 if sample_weights is not None:
330 grad_values = grad_values * sample_weights
331 grad_input = grad_input.contiguous()
332 grad_input[row_indices[in_vocab_mask], local_targets[in_vocab_mask]] += grad_values[in_vocab_mask]
334 if not ignore_mask.all():
335 if state.reduction == "none":
336 grad_input[~ignore_mask] = 0.0
337 else:
338 ignore_indices_expanded = (~ignore_mask).unsqueeze(-1).expand_as(grad_input)
339 grad_input[ignore_indices_expanded] = 0.0
340 return grad_input
343class DistributedCrossEntropyFunction(torch.autograd.Function):
344 """K3: Fused backward for distributed cross_entropy."""
346 @staticmethod
347 def forward(
348 ctx: Any,
349 input_local: Tensor,
350 target: Tensor,
351 weight: Optional[Tensor],
352 ignore_index: int,
353 reduction: str,
354 vocab_size: int,
355 mesh: DeviceMesh,
356 mesh_dim: int,
357 ) -> Tensor:
358 """Forward pass with shared context setup for each reduction mode."""
359 ctx.vocab_start = _compute_vocab_start(
360 vocab_size, mesh.size(mesh_dim), mesh.get_local_rank(mesh_dim)
361 )
362 ctx.vocab_end = ctx.vocab_start + input_local.shape[-1]
363 ctx.reduction = reduction
364 ctx.ignore_index = ignore_index
365 log_probs_local = distributed_log_softmax(
366 input_local, dim=-1, mesh=mesh, mesh_dim=mesh_dim
367 )
368 nll_result = distributed_nll_loss_forward(
369 log_probs_local,
370 target,
371 weight,
372 ignore_index,
373 reduction,
374 ctx.vocab_start,
375 ctx.vocab_end,
376 )
377 loss, total_weight = nll_result[:2]
378 if reduction in ("mean", "sum"):
379 group = mesh.get_group(mesh_dim)
380 loss = _differentiable_all_reduce(loss, op="sum", group=group)
381 if reduction == "mean":
382 total_weight = _differentiable_all_reduce(
383 total_weight, op="sum", group=group
384 )
385 if reduction != "mean":
386 total_weight = torch.zeros(1, dtype=loss.dtype, device=loss.device)
387 ctx.save_for_backward(
388 log_probs_local, target, weight, total_weight, nll_result[2], nll_result[3]
389 )
390 if reduction == "mean":
391 if total_weight.item() == 0:
392 return torch.tensor(float('nan'), dtype=loss.dtype, device=loss.device)
393 return loss / total_weight
394 return loss
396 @staticmethod
397 def backward(ctx: Any, grad_output: Tensor) -> Tuple[Optional[Tensor], ...]:
398 """Backward pass (vectorized implementation)."""
399 state = _CrossEntropyContext(
400 *ctx.saved_tensors,
401 reduction=ctx.reduction,
402 ignore_index=ctx.ignore_index,
403 vocab_start=ctx.vocab_start,
404 vocab_end=ctx.vocab_end,
405 )
406 return _CrossEntropyGradients(_compute_cross_entropy_gradient(state, grad_output))
409def vocab_parallel_cross_entropy_local(
410 local_logits: Tensor,
411 target: Tensor,
412 *,
413 vocab_size: int,
414 mesh: "DeviceMesh",
415 mesh_dim: Optional[int] = None,
416 ignore_index: int = -100,
417 reduction: str = "sum",
418) -> Tensor:
419 """Vocab-parallel cross entropy on a local logits shard (local-first).
421 AutoModels production entry: the caller owns the sharding contract and
422 passes the local shard directly, so no temporary DTensor is created per
423 call. ``mesh``/``mesh_dim`` identify the single mesh axis that shards the
424 class dimension; see :func:`_resolve_class_mesh_dim`.
426 Args:
427 local_logits: Local class-sharded logits, ``[..., local_vocab_size]``.
428 Rank above 2 is flattened to ``[-1, local_vocab_size]`` together
429 with ``target``.
430 target: Class indices aligned with the leading dims of
431 ``local_logits``.
432 vocab_size: Global vocabulary size (across all shard ranks).
433 mesh: Mesh holding the class-sharding axis.
434 mesh_dim: Explicit class-sharding mesh axis; ``None`` resolves it
435 from the mesh (1D mesh → axis 0; otherwise the axis named "tp").
436 ignore_index: Target value excluded from the loss.
437 reduction: Reduction method: 'none', 'mean', or 'sum'.
439 Returns:
440 Loss tensor.
442 Raises:
443 ValueError: If the class-sharding mesh axis cannot be uniquely
444 identified, the target is probabilistic, or the reduction is
445 invalid.
446 """
447 mesh_dim = _resolve_class_mesh_dim(mesh, mesh_dim)
448 _validate_target_type_base(_is_floating_torch(target))
449 if reduction not in ("none", "mean", "sum"):
450 raise ValueError(f"Invalid reduction: {reduction}. Must be 'none', 'mean', or 'sum'.")
452 local_vocab_size = local_logits.shape[-1]
453 if local_logits.ndim > 2:
454 local_logits = local_logits.reshape(-1, local_vocab_size)
455 target = target.reshape(-1)
457 return DistributedCrossEntropyFunction.apply(
458 local_logits,
459 target,
460 None,
461 ignore_index,
462 reduction,
463 vocab_size,
464 mesh,
465 mesh_dim,
466 )