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

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 calculation dispatcher — calculate_loss following design doc §10.""" 

16 

17__all__ = ["calculate_loss"] 

18 

19from typing import Any, Optional 

20 

21import torch 

22from torch import nn 

23 

24# FusedLinearCrossEntropy is a stub placeholder — actual implementation 

25# depends on cut_cross_entropy or similar fused kernel. The dispatcher 

26# handles it via type check; if missing, all paths fall through to logit-based CE. 

27try: 

28 from hyper_parallel.components.losses.linear_ce import FusedLinearCrossEntropy 

29except ImportError: 

30 FusedLinearCrossEntropy = None 

31 

32 

33def _get_lm_head_weight(model: nn.Module) -> Optional[torch.Tensor]: 

34 """Return the model's LM-head weight, materializing DTensor when needed.""" 

35 if hasattr(model, "get_output_embeddings"): 

36 lm_head = model.get_output_embeddings() 

37 if lm_head is not None and hasattr(lm_head, "weight"): 

38 weight = lm_head.weight 

39 return weight.full_tensor() if hasattr(weight, "full_tensor") else weight 

40 for name, param in model.named_parameters(remove_duplicate=False): 

41 if "lm_head" in name and name.endswith(".weight"): 

42 return param.full_tensor() if hasattr(param, "full_tensor") else param 

43 return None 

44 

45 

46def calculate_loss(loss_fn: nn.Module, **kwargs: Any) -> torch.Tensor: 

47 """Unified loss calculation dispatcher. 

48 

49 Follows design doc §10: 

50 - Path A: FusedLinearCrossEntropy — uses hidden_states + lm_weight 

51 - Path B: Standard logit-based CE — uses logits + labels 

52 

53 Returns raw ce_sum (not divided by N) for token_weighted aggregation. 

54 Token-mean normalization is done by scale_grads_and_clip_grad_norm. 

55 

56 Args: 

57 loss_fn: Loss module (MaskedCrossEntropy, FusedLinearCrossEntropy, etc.) 

58 **kwargs: Must contain logits+labels (Path B) or hidden_states+lm_weight (Path A). 

59 

60 Returns: 

61 Raw loss tensor (reduction="sum", not divided by num_label_tokens). 

62 """ 

63 loss_aggregation = kwargs.pop("loss_aggregation", "token_weighted") 

64 

65 if FusedLinearCrossEntropy is not None and isinstance(loss_fn, FusedLinearCrossEntropy): 

66 # ── Path A: FusedLinearCrossEntropy ── 

67 hidden_states = kwargs.get("hidden_states") 

68 lm_weight = kwargs.get("lm_weight") 

69 labels = kwargs.get("labels") 

70 if hidden_states is not None and lm_weight is not None: 

71 return loss_fn( 

72 hidden_states=hidden_states, 

73 labels=labels, 

74 lm_weight=lm_weight, 

75 ) 

76 # Fallback: use logits 

77 logits = kwargs.get("logits") 

78 if logits is None: 

79 raise ValueError( 

80 "FusedLinearCrossEntropy requires hidden_states+lm_weight or logits" 

81 ) 

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

83 labels = labels[..., 1:].contiguous() 

84 return loss_fn( 

85 logits.view(-1, logits.size(-1)), 

86 labels.view(-1), 

87 ) 

88 

89 # ── Path B: Standard logit-based loss ── 

90 logits = kwargs["logits"] 

91 labels = kwargs["labels"] 

92 

93 # Shift: causal LM always needs shift 

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

95 labels = labels[..., 1:].contiguous() 

96 

97 if loss_aggregation == "token_weighted": 

98 # Return raw ce_sum; token-mean normalization by scale_grads 

99 return loss_fn( 

100 logits.view(-1, logits.size(-1)), 

101 labels.view(-1), 

102 ) 

103 # rank_average: mean-scale loss 

104 return loss_fn( 

105 logits.view(-1, logits.size(-1)), 

106 labels.view(-1), 

107 )