Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / losses / chunked_cross_entropy.py: 84%
156 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"""Memory-efficient chunked output projection and cross-entropy.
17The full ``[batch, sequence, vocabulary]`` logits tensor is never
18materialized. Output projection, FP32 cross-entropy, and first-order
19gradients are evaluated one local sequence chunk at a time.
20"""
22from __future__ import annotations
24__all__ = [
25 "ChunkedCausalLMLoss",
26 "ChunkedCausalLMOutput",
27 "chunked_cross_entropy",
28]
30from collections.abc import Mapping
31from dataclasses import dataclass
32from typing import Any
34# This loss is a PyTorch Trainer component and differentiates through
35# torch.func; there is no platform-neutral autograd equivalent. PyTorch's
36# Function metaclass supplies its public variadic dispatch contract.
37# pylint: disable=forbidden-backend-import,not-callable,abstract-method,arguments-differ
38import torch
39from torch import nn
40from torch.autograd.function import once_differentiable
41from torch.nn import functional
42from transformers.utils import ModelOutput
45def _linear_cross_entropy_chunk(
46 hidden_chunk: torch.Tensor,
47 weight: torch.Tensor,
48 target_chunk: torch.Tensor,
49 ignore_index: int,
50) -> torch.Tensor:
51 """Return the FP32 summed cross-entropy for one sequence chunk."""
52 logits = functional.linear(
53 hidden_chunk.reshape(-1, hidden_chunk.size(-1)),
54 weight,
55 ).float()
56 return functional.cross_entropy(
57 logits,
58 target_chunk.reshape(-1),
59 ignore_index=ignore_index,
60 reduction="sum",
61 )
64class _PrecomputedChunkLoss(torch.autograd.Function):
65 """Precompute per-chunk gradients and replay them in outer backward."""
67 @staticmethod
68 def forward(
69 ctx: Any,
70 hidden_states: torch.Tensor,
71 head_weight: torch.Tensor,
72 targets: torch.Tensor,
73 chunk_size: int,
74 ignore_index: int,
75 ) -> torch.Tensor:
76 """Compute the loss sum and retain accumulated first-order gradients."""
77 grad_hidden = torch.empty_like(hidden_states)
78 grad_weight = torch.zeros_like(head_weight)
79 loss_sum = torch.zeros((), dtype=torch.float32, device=hidden_states.device)
81 grad_and_value = torch.func.grad_and_value(
82 _linear_cross_entropy_chunk,
83 argnums=(0, 1),
84 )
85 hidden_chunks = torch.split(hidden_states, chunk_size, dim=1)
86 target_chunks = torch.split(targets, chunk_size, dim=1)
87 grad_hidden_chunks = torch.split(grad_hidden, chunk_size, dim=1)
88 for hidden_chunk, target_chunk, grad_hidden_chunk in zip(
89 hidden_chunks, target_chunks, grad_hidden_chunks, strict=True
90 ):
91 (chunk_grad_hidden, chunk_grad_weight), chunk_loss = grad_and_value(
92 hidden_chunk,
93 head_weight,
94 target_chunk,
95 ignore_index,
96 )
97 grad_hidden_chunk.copy_(chunk_grad_hidden)
98 grad_weight.add_(chunk_grad_weight)
99 loss_sum.add_(chunk_loss)
100 del chunk_grad_hidden, chunk_grad_weight, chunk_loss
102 ctx.save_for_backward(grad_hidden, grad_weight)
103 return loss_sum
105 @staticmethod
106 @once_differentiable
107 def backward(ctx: Any, grad_loss_sum: torch.Tensor | None) -> tuple:
108 """Replay private gradients with the Trainer's upstream scale."""
109 if grad_loss_sum is None:
110 return None, None, None, None, None
111 grad_hidden, grad_weight = ctx.saved_tensors
112 return (
113 grad_hidden * grad_loss_sum,
114 grad_weight * grad_loss_sum,
115 None,
116 None,
117 None,
118 )
121def _validate_chunk_options(chunk_size: int, ignore_index: int) -> None:
122 """Validate scalar chunk-loss options."""
123 if isinstance(chunk_size, bool) or not isinstance(chunk_size, int):
124 raise TypeError(
125 f"chunk_size must be an integer, got {type(chunk_size).__name__}"
126 )
127 if chunk_size <= 0:
128 raise ValueError(f"chunk_size must be positive, got {chunk_size}")
129 if isinstance(ignore_index, bool) or not isinstance(ignore_index, int):
130 raise TypeError(
131 f"ignore_index must be an integer, got {type(ignore_index).__name__}"
132 )
135def _validate_chunk_shapes(
136 hidden_states: torch.Tensor,
137 targets: torch.Tensor,
138 head_weight: torch.Tensor,
139) -> None:
140 """Validate tensor ranks and logical dimensions."""
141 if hidden_states.dim() != 3:
142 raise ValueError(
143 "hidden_states must have shape [batch, sequence, hidden], "
144 f"got {tuple(hidden_states.shape)}"
145 )
146 if targets.dim() != 2 or hidden_states.shape[:2] != targets.shape:
147 raise ValueError(
148 "targets must match hidden_states batch/sequence dimensions, "
149 f"got hidden={tuple(hidden_states.shape)} and targets={tuple(targets.shape)}"
150 )
151 if hidden_states.size(1) == 0 or hidden_states.size(2) == 0:
152 raise ValueError(
153 "hidden_states sequence and hidden dimensions must be non-empty"
154 )
155 if head_weight.dim() != 2 or head_weight.size(1) != hidden_states.size(2):
156 raise ValueError(
157 "head_weight must have shape [vocabulary, hidden] matching "
158 f"hidden_states, got hidden={tuple(hidden_states.shape)} and "
159 f"weight={tuple(head_weight.shape)}"
160 )
161 if head_weight.size(0) == 0:
162 raise ValueError("head_weight vocabulary dimension must be non-empty")
165def _validate_chunk_dtypes_and_device(
166 hidden_states: torch.Tensor,
167 targets: torch.Tensor,
168 head_weight: torch.Tensor,
169) -> None:
170 """Validate tensor dtypes and device ownership."""
171 if not torch.is_floating_point(hidden_states) or not torch.is_floating_point(
172 head_weight
173 ):
174 raise TypeError("hidden_states and head_weight must be floating-point tensors")
175 if targets.dtype != torch.long:
176 raise TypeError(f"targets must have dtype torch.long, got {targets.dtype}")
177 if hidden_states.dtype != head_weight.dtype:
178 raise TypeError(
179 "hidden_states and head_weight must have the same dtype, "
180 f"got {hidden_states.dtype} and {head_weight.dtype}"
181 )
182 if not hidden_states.device == targets.device == head_weight.device:
183 raise ValueError(
184 "hidden_states, targets, and head_weight must be on the same device, "
185 f"got {hidden_states.device}, {targets.device}, and {head_weight.device}"
186 )
189def chunked_cross_entropy(
190 hidden_states: torch.Tensor,
191 targets: torch.Tensor,
192 head_weight: torch.Tensor,
193 chunk_size: int = 1024,
194 ignore_index: int = -100,
195) -> torch.Tensor:
196 """Compute summed linear cross-entropy without full-sequence logits.
198 Args:
199 hidden_states: Local hidden states in ``[B, S, H]`` layout.
200 targets: Integer targets aligned with ``hidden_states`` in ``[B, S]``.
201 head_weight: Full-vocabulary output weight in ``[V, H]`` layout.
202 chunk_size: Maximum local sequence length evaluated per chunk.
203 ignore_index: Target value excluded from the summed loss.
205 Returns:
206 A scalar FP32 loss sum connected to ``hidden_states`` and
207 ``head_weight``.
208 """
209 _validate_chunk_options(chunk_size, ignore_index)
210 _validate_chunk_shapes(hidden_states, targets, head_weight)
211 _validate_chunk_dtypes_and_device(hidden_states, targets, head_weight)
212 return _PrecomputedChunkLoss.apply(
213 hidden_states,
214 head_weight,
215 targets,
216 chunk_size,
217 ignore_index,
218 )
221def _get_output_value(model_output: Any, name: str) -> Any:
222 """Read one optional value from a mapping or Transformers output."""
223 if isinstance(model_output, Mapping):
224 return model_output.get(name)
225 return getattr(model_output, name, None)
228@dataclass
229class ChunkedCausalLMOutput(ModelOutput):
230 """Training output produced by a model-integrated Chunk Loss adapter."""
232 loss_sum: torch.Tensor | None = None
233 valid_token_count: torch.Tensor | None = None
234 aux_loss: torch.Tensor | None = None
235 aux_loss_coef: float | None = None
236 logits: torch.Tensor | None = None
237 past_key_values: Any | None = None
238 hidden_states: Any | None = None
239 attentions: Any | None = None
240 router_logits: Any | None = None
243class ChunkedCausalLMLoss(nn.Module):
244 """Trainer-facing model-integrated Chunk Loss objective.
246 The configured model-family adapter consumes the prepared model inputs,
247 calls :func:`chunked_cross_entropy` before a full LM-head output exists,
248 and returns :class:`ChunkedCausalLMOutput`. This module converts the local
249 summed CE to one local token mean; global DP/CP weighting remains owned by
250 ``mean_global_loss``.
251 """
253 def __init__(self, chunk_size: int = 1024, ignore_index: int = -100) -> None:
254 """Initialize and validate the per-rank sequence chunk options."""
255 super().__init__()
256 _validate_chunk_options(chunk_size, ignore_index)
257 self.chunk_size = chunk_size
258 self.ignore_index = ignore_index
259 self._model_is_bound = False
261 @staticmethod
262 def _parallel_size(distributed_setup: Any, name: str) -> int:
263 """Read one topology size from the shared mesh context."""
264 mesh = getattr(distributed_setup, "mesh_context", None)
265 return int(getattr(mesh, name, 1))
267 def bind_model(self, model: nn.Module, distributed_setup: Any = None) -> None:
268 """Bind the registered model-family Chunk Loss forward adapter.
270 Args:
271 model: Fully constructed model whose parameter structure is fixed.
272 distributed_setup: Shared Trainer distributed setup.
274 Raises:
275 NotImplementedError: If an unsupported TP/PP/loss-parallel mode is
276 active.
277 ValueError: If the model family has no Chunk Loss adapter.
278 """
279 # Trainer-side token aggregation currently follows the data contract's
280 # canonical ignore value. The model-integrated objective must use the
281 # same value or DP/CP weighting would count a different token set.
282 from hyper_parallel.data.constants import IGNORE_INDEX # pylint: disable=import-outside-toplevel
284 if self.ignore_index != IGNORE_INDEX:
285 raise ValueError(
286 "ChunkedCausalLMLoss Trainer integration requires "
287 f"ignore_index={IGNORE_INDEX}, got {self.ignore_index}"
288 )
290 tp_size = self._parallel_size(distributed_setup, "tp_size")
291 pp_size = self._parallel_size(distributed_setup, "pp_size")
292 loss_parallel = bool(
293 getattr(getattr(distributed_setup, "mesh_context", None), "loss_parallel", False)
294 )
295 if tp_size != 1 or pp_size != 1 or loss_parallel:
296 raise NotImplementedError(
297 "ChunkedCausalLMLoss currently requires tp_size=1, pp_size=1, "
298 f"and loss_parallel=false; got tp_size={tp_size}, "
299 f"pp_size={pp_size}, loss_parallel={loss_parallel}"
300 )
302 # Lazy imports keep generic loss-package import independent of concrete
303 # model families and their optional backend dependencies.
304 from hyper_parallel.models.registry import get_model_adapter # pylint: disable=import-outside-toplevel
306 model_type = getattr(getattr(model, "config", None), "model_type", "")
307 adapter_spec = get_model_adapter(model_type)
308 if adapter_spec is None:
309 architectures = getattr(model.config, "architectures", None) or ()
310 if architectures:
311 adapter_spec = get_model_adapter(architectures[0])
312 provider = None if adapter_spec is None else adapter_spec.loss
313 if provider is None:
314 raise ValueError(
315 "ChunkedCausalLMLoss requires a registered model-family loss "
316 f"adapter, but model_type={model_type!r} and its architecture provide none"
317 )
318 adapter = provider()
319 bind_chunk_loss = getattr(adapter, "bind_chunk_loss", None)
320 if not callable(bind_chunk_loss):
321 raise TypeError(
322 f"model_type={model_type!r} loss adapter must expose bind_chunk_loss(model)"
323 )
324 bind_chunk_loss(model)
325 self._model_is_bound = True
327 def prepare_model_inputs(
328 self,
329 model_inputs: Mapping[str, Any],
330 loss_inputs: Mapping[str, Any],
331 ) -> dict[str, Any]:
332 """Move labels to the model-family no-full-logits protocol."""
333 if not self._model_is_bound:
334 raise RuntimeError("ChunkedCausalLMLoss must be bound before model forward")
335 targets = loss_inputs.get("shift_labels")
336 if not isinstance(targets, torch.Tensor):
337 raise ValueError(
338 "ChunkedCausalLMLoss requires pre-shifted shift_labels aligned "
339 "with every local hidden position"
340 )
341 loss_mask = loss_inputs.get("loss_mask")
342 if loss_mask is not None:
343 if not isinstance(loss_mask, torch.Tensor) or loss_mask.shape != targets.shape:
344 raise ValueError("loss_mask must be a Tensor with the same shape as targets")
346 prepared = dict(model_inputs)
347 # A legacy one-dictionary caller may leave loss-only fields in
348 # model_inputs. Replace them with the adapter's explicit protocol so
349 # labels and masks never leak into the decoder's generic **kwargs.
350 prepared.pop("labels", None)
351 prepared.pop("shift_labels", None)
352 prepared.pop("loss_mask", None)
353 prepared.update(
354 {
355 "chunk_loss_targets": targets,
356 "chunk_loss_mask": loss_mask,
357 "chunk_loss_chunk_size": self.chunk_size,
358 "chunk_loss_ignore_index": self.ignore_index,
359 }
360 )
361 return prepared
363 def forward( # pylint: disable=unused-argument
364 self,
365 *,
366 model_output: Any,
367 labels: torch.Tensor | None,
368 ) -> torch.Tensor:
369 """Normalize one adapter-produced CE sum exactly once."""
370 logits = _get_output_value(model_output, "logits")
371 if logits is not None:
372 raise ValueError(
373 "ChunkedCausalLMLoss model adapter materialized logits; expected logits=None"
374 )
375 loss_sum = _get_output_value(model_output, "loss_sum")
376 valid_token_count = _get_output_value(model_output, "valid_token_count")
377 if not isinstance(loss_sum, torch.Tensor) or loss_sum.numel() != 1:
378 raise ValueError("ChunkedCausalLMLoss requires a scalar Tensor loss_sum")
379 if not isinstance(valid_token_count, torch.Tensor) or valid_token_count.numel() != 1:
380 raise ValueError(
381 "ChunkedCausalLMLoss requires a scalar Tensor valid_token_count"
382 )
384 denominator = valid_token_count.clamp_min(1).to(loss_sum.dtype)
385 local_loss = loss_sum / denominator
386 aux_loss = _get_output_value(model_output, "aux_loss")
387 if aux_loss is not None:
388 if not isinstance(aux_loss, torch.Tensor) or aux_loss.numel() != 1:
389 raise ValueError("ChunkedCausalLMLoss aux_loss must be a scalar Tensor")
390 aux_loss_coef = _get_output_value(model_output, "aux_loss_coef")
391 local_loss = local_loss + float(
392 0.0 if aux_loss_coef is None else aux_loss_coef
393 ) * aux_loss
394 return torch.where(
395 valid_token_count.ne(0),
396 local_loss,
397 torch.zeros_like(local_loss),
398 )