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

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"] 

18 

19import torch 

20from torch import nn 

21 

22 

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

24 mtp_per_depth_logits: list[torch.Tensor], 

25 mtp_per_depth_h: list[torch.Tensor], 

26 labels: torch.Tensor, 

27 loss_fn: nn.Module, 

28) -> torch.Tensor: 

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

30 

31 Computes CE per depth and sums them. 

32 

33 Args: 

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

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

36 variants that condition the loss on hidden states; currently 

37 unused. 

38 labels: Target token indices. 

39 loss_fn: Loss module applied per depth. 

40 

41 Returns: 

42 Summed MTP loss over all depths. 

43 """ 

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

45 for logits in mtp_per_depth_logits: 

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

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

48 depth_loss = loss_fn( 

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

50 labels_shifted.view(-1), 

51 ) 

52 total_mtp_loss = total_mtp_loss + depth_loss 

53 return total_mtp_loss