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

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"""Loss module that reads the loss produced by a model.""" 

16 

17__all__ = ["ModelOutputLoss"] 

18 

19from typing import Any, Dict, Optional, Union 

20 

21# AutoModels loss components implement the Transformers/PyTorch Trainer API. 

22# pylint: disable-next=forbidden-backend-import 

23import torch 

24 

25from hyper_parallel.data.constants import IGNORE_INDEX 

26 

27 

28class ModelOutputLoss(torch.nn.Module): 

29 """Return the loss field from a Transformers-style model output.""" 

30 

31 def forward( # pylint: disable=unused-argument 

32 self, 

33 *, 

34 model_output: Any, 

35 labels: Optional[torch.Tensor], 

36 ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: 

37 """Read the model-computed loss. 

38 

39 Args: 

40 model_output: Model output exposing a ``loss`` attribute. 

41 labels: Labels associated with the output. This default loss keeps 

42 the argument only to share the trainer-facing call signature 

43 with replaceable loss modules. 

44 

45 Returns: 

46 The loss tensor or named loss mapping from ``model_output.loss``. 

47 """ 

48 local_loss = model_output.loss 

49 if labels is None or not isinstance(local_loss, torch.Tensor): 

50 return local_loss 

51 

52 # Causal LM loss shifts labels by one position. A CP-local slice may 

53 # therefore contain no trainable target even when other CP ranks do. 

54 has_valid_labels = labels[..., 1:].ne(IGNORE_INDEX).any() 

55 local_loss = torch.where(has_valid_labels, local_loss, torch.zeros_like(local_loss)) 

56 

57 return local_loss