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

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"""FusedLinearCrossEntropy — stub for fused lm_head + CE loss. 

16 

17Reference: nemo_automodel.components.loss.linear_ce 

18Full implementation requires cut_cross_entropy package. 

19""" 

20 

21from typing import Optional 

22 

23import torch 

24from torch import nn 

25import torch.nn.functional as F 

26 

27 

28class FusedLinearCrossEntropy(nn.Module): 

29 """Fused linear + cross-entropy loss. 

30 

31 Stub implementation: falls back to separate linear + CE. 

32 The production version fuses lm_head matmul with CE computation 

33 for memory efficiency (avoids materializing the full logits tensor). 

34 """ 

35 

36 def __init__(self, fp32_upcast: bool = True, ignore_index: int = -100) -> None: 

37 """Initialize the fused linear + cross-entropy loss module. 

38 

39 Args: 

40 fp32_upcast: Whether to upcast inputs to float32 before the 

41 matmul / CE computation for numerical stability. 

42 ignore_index: Label index ignored by the cross-entropy loss. 

43 """ 

44 super().__init__() 

45 self.fp32_upcast = fp32_upcast 

46 self.ignore_index = ignore_index 

47 

48 def forward( 

49 self, 

50 hidden_states: torch.Tensor, 

51 labels: torch.Tensor, 

52 lm_weight: Optional[torch.Tensor] = None, 

53 num_label_tokens: Optional[int] = None, 

54 ) -> torch.Tensor: 

55 """Compute fused linear + cross-entropy loss. 

56 

57 Args: 

58 hidden_states: [B, S, H] final hidden states (when ``lm_weight`` 

59 is provided), or pre-computed logits [N, V] (when 

60 ``lm_weight=None``, i.e. the logits fallback path of the 

61 ``calculate_loss`` dispatcher). 

62 labels: [B, S] or [N] target indices. 

63 lm_weight: [V, H] LM head weight; when None, ``hidden_states`` is 

64 treated as pre-computed logits. 

65 num_label_tokens: Optional token count for normalization. 

66 

67 Returns: 

68 Scalar loss (reduction="sum"). 

69 """ 

70 if lm_weight is not None: 

71 logits = torch.matmul(hidden_states.float() if self.fp32_upcast else hidden_states, lm_weight.t()) 

72 else: 

73 # dispatcher fallback path: hidden_states are pre-computed logits 

74 logits = hidden_states.float() if self.fp32_upcast else hidden_states 

75 logits = logits.view(-1, logits.size(-1)) 

76 labels = labels.view(-1) 

77 loss = F.cross_entropy(logits, labels, reduction="sum", ignore_index=self.ignore_index) 

78 if num_label_tokens is not None: 

79 if num_label_tokens > 0: 

80 loss = loss / num_label_tokens 

81 return loss