Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / losses / mtp.py: 97%

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"""Multi-Token-Prediction auxiliary loss objective.""" 

16 

17__all__ = ["calculate_mtp_loss", "compute_mtp_loss"] 

18 

19import math 

20from collections.abc import Sequence 

21from typing import Literal 

22 

23import torch 

24from torch import nn 

25from torch.nn import functional as F 

26 

27 

28def calculate_mtp_loss( # pylint: disable=unused-argument 

29 mtp_per_depth_logits: list[torch.Tensor], 

30 mtp_per_depth_h: list[torch.Tensor], 

31 labels: torch.Tensor, 

32 loss_fn: nn.Module, 

33) -> torch.Tensor: 

34 """Multi-Token-Prediction auxiliary loss. 

35 

36 Computes CE per depth and sums them. 

37 

38 Args: 

39 mtp_per_depth_logits: Per-depth logits from the MTP heads. 

40 mtp_per_depth_h: Per-depth hidden states. Reserved for future MTP 

41 variants that condition the loss on hidden states; currently 

42 unused. 

43 labels: Target token indices. 

44 loss_fn: Loss module applied per depth. 

45 

46 Returns: 

47 Summed MTP loss over all depths. 

48 """ 

49 total_mtp_loss = torch.tensor(0.0, device=labels.device, dtype=torch.float32) 

50 for logits in mtp_per_depth_logits: 

51 logits_shifted = logits[..., :-1, :].contiguous() 

52 labels_shifted = labels[..., 1:].contiguous() 

53 depth_loss = loss_fn( 

54 logits_shifted.view(-1, logits_shifted.size(-1)), 

55 labels_shifted.view(-1), 

56 ) 

57 total_mtp_loss = total_mtp_loss + depth_loss 

58 return total_mtp_loss 

59 

60 

61def compute_mtp_loss( 

62 per_depth_logits: Sequence[torch.Tensor], 

63 next_token_labels: torch.Tensor, 

64 *, 

65 loss_factor: float = 0.3, 

66 normalization: Literal["sequence", "valid_tokens"] = "sequence", 

67 ignore_index: int = -100, 

68) -> torch.Tensor: 

69 """Compute loss_factor / D times the sum of per-depth cross-entropies. 

70 

71 Args: 

72 per_depth_logits: Full-vocabulary logits, depth k shaped [B, S-k, V]. 

73 next_token_labels: ALREADY shifted main-LM targets [B, S]. If input_ids 

74 are [t0, t1, t2, t3], these labels are [t1, t2, t3, t4]; if t4 is 

75 unavailable, use ignore_index. Depth 1 scores [t2, t3, t4]. 

76 loss_factor: Nonnegative finite weight for the MTP objective. 

77 normalization: 'sequence' divides each depth's loss sum by B*S, matching 

78 the original-T denominator in DeepSeek-V3 paper Eq. 24. 'valid_tokens' 

79 divides by each depth's unmasked target count. The choice changes 

80 gradient scaling and is explicit; neither is a claim about unpublished 

81 DeepSeek training code. Empty valid-target sets contribute zero. 

82 ignore_index: Label excluded from cross-entropy. 

83 

84 Returns: 

85 Scalar FP32 weighted MTP loss, connected to every depth's logits. 

86 

87 Note: 

88 No DP/CP/global accumulation normalization or vocabulary-sharded CE is 

89 performed here. Distributed adapters own those reductions. The existing 

90 calculate_mtp_loss API is intentionally unchanged. 

91 """ 

92 if not per_depth_logits: 

93 raise ValueError("per_depth_logits must contain at least one prediction depth") 

94 if next_token_labels.ndim != 2 or next_token_labels.numel() == 0: 

95 raise ValueError("next_token_labels must be nonempty [B, S]") 

96 if next_token_labels.dtype != torch.long: 

97 raise ValueError("next_token_labels must have dtype torch.long") 

98 if normalization not in ("sequence", "valid_tokens"): 

99 raise ValueError("normalization must be 'sequence' or 'valid_tokens'") 

100 if not math.isfinite(loss_factor) or loss_factor < 0: 

101 raise ValueError("loss_factor must be finite and nonnegative") 

102 total = torch.zeros((), device=next_token_labels.device, dtype=torch.float32) 

103 for depth, logits in enumerate(per_depth_logits, start=1): 

104 targets = next_token_labels[:, depth:] 

105 if targets.shape[1] == 0 or logits.ndim != 3 or logits.shape[:2] != targets.shape: 

106 raise ValueError("Depth k logits must have shape [B, S-k, V] for next_token_labels [B, S]") 

107 loss_sum = F.cross_entropy(logits.float().reshape(-1, logits.shape[-1]), targets.reshape(-1), 

108 ignore_index=ignore_index, reduction="sum") 

109 denominator = (next_token_labels.numel() if normalization == "sequence" 

110 else (targets != ignore_index).sum().clamp_min(1)) 

111 total = total + loss_sum / denominator 

112 return total * (loss_factor / len(per_depth_logits))