Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / runtime / loss_aggregation.py: 67%

1# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates 

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"""Loss token counting for trainer-side aggregation.""" 

15 

16from typing import Union 

17 

18import torch 

19 

20from hyper_parallel.data.constants import IGNORE_INDEX 

21 

22 

23def count_loss_token( 

24 batches: Union[list[dict[str, torch.Tensor]], dict[str, torch.Tensor]], 

25) -> dict[str, torch.Tensor]: 

26 # FIXME: VeOmni version -> AutoModels version 

27 """Calculate the total number of text_tokens/image_tokens/** for loss in a global batch, or one micro batch.""" 

28 if isinstance(batches, dict): 

29 batches = [batches] 

30 token_len: dict[str, torch.Tensor] = {} 

31 

32 def _count(obj): 

33 if isinstance(obj, dict) and not obj.get("padding_flag", False): 

34 # ``shift_labels`` already aligns one target with every local hidden 

35 # position. Otherwise align both labels and mask with the causal 

36 # next-token prediction before counting. 

37 labels = obj.get("shift_labels") 

38 loss_mask = obj.get("loss_mask") 

39 if labels is None: 

40 labels = obj["labels"][..., 1:] 

41 if loss_mask is not None: 

42 loss_mask = loss_mask[..., 1:] 

43 # IGNORE_INDEX and loss_mask are independent exclusion mechanisms; 

44 # global loss weighting must count the same intersection as CE. 

45 valid_targets = labels != IGNORE_INDEX 

46 if loss_mask is not None: 

47 if loss_mask.shape != labels.shape: 

48 raise ValueError( 

49 "loss_mask must match the labels after causal alignment, " 

50 f"got labels={tuple(labels.shape)} and " 

51 f"aligned_mask={tuple(loss_mask.shape)}" 

52 ) 

53 valid_targets = valid_targets & loss_mask.to(torch.bool) 

54 foundation_tokens = torch.sum(valid_targets) 

55 if "foundation_tokens" in token_len: 

56 foundation_tokens = token_len["foundation_tokens"] + foundation_tokens 

57 token_len["foundation_tokens"] = foundation_tokens # text tokens 

58 

59 for key in obj.keys(): 

60 if key.endswith("_labels") and key != "shift_labels": 

61 token_name = key.split("_labels")[0] 

62 token_len[f"{token_name}_tokens"] = torch.sum(obj[key] != IGNORE_INDEX) # image generation tokens 

63 

64 if "image_output_mask" in obj: 

65 image_decoder_tokens = torch.sum(obj["image_output_mask"]) 

66 if "image_decoder_tokens" in token_len: 

67 image_decoder_tokens = token_len["image_decoder_tokens"] + image_decoder_tokens 

68 token_len["image_decoder_tokens"] = image_decoder_tokens # image generation tokens 

69 elif isinstance(obj, (list, tuple)): 

70 for item in obj: 

71 _count(item) 

72 else: 

73 raise TypeError(f"Unsupported batch type: {type(obj)}") 

74 

75 _count(batches) 

76 foundation_tokens = token_len.setdefault("foundation_tokens", torch.tensor(0)) 

77 token_len.setdefault("image_decoder_tokens", foundation_tokens.new_zeros(())) 

78 return token_len 

79 

80 

81__all__ = ["count_loss_token"]