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

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

2# Copyright 2026 Huawei Technologies Co., Ltd 

3# 

4# Licensed under the Apache License, Version 2.0 (the "License"); 

5# you may not use this file except in compliance with the License. 

6# You may obtain a copy of the License at 

7# 

8# http://www.apache.org/licenses/LICENSE-2.0 

9# 

10# Unless required by applicable law or agreed to in writing, software 

11# distributed under the License is distributed on an "AS IS" BASIS, 

12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

13# See the License for the specific language governing permissions and 

14# limitations under the License. 

15# ============================================================================ 

16"""Global loss metrics reduction across data/context parallel ranks.""" 

17 

18from __future__ import annotations 

19 

20__all__ = ["mean_global_loss"] 

21 

22from typing import TYPE_CHECKING, Union 

23 

24import torch # pylint: disable=forbidden-backend-import 

25import torch.distributed as dist # pylint: disable=forbidden-backend-import 

26 

27from hyper_parallel.core.utils.communication import ( 

28 _is_tracing_or_compiling, 

29 differentiable_all_reduce, 

30) 

31from hyper_parallel.trainer.runtime.distributed import all_reduce 

32 

33if TYPE_CHECKING: 

34 from hyper_parallel.distributed.mesh import MeshContext 

35 

36 

37def _reduce_weighted_loss(cur_loss, cur_token_len, device_mesh, dp_cp_group): 

38 """Reduce a loss numerator while preserving the local backward graph.""" 

39 local_weighted_loss = cur_loss.float() * cur_token_len 

40 backward_loss = local_weighted_loss * device_mesh.dp_size * device_mesh.cp_size 

41 global_weighted_loss = local_weighted_loss.detach().item() 

42 if device_mesh.dp_cp_mesh is not None and device_mesh.dp_cp_mesh.size() > 1: 

43 global_weighted_loss = all_reduce( 

44 global_weighted_loss, 

45 op="sum", 

46 group=dp_cp_group, 

47 ) 

48 global_sum = local_weighted_loss.new_tensor(global_weighted_loss) 

49 return global_sum + (backward_loss - backward_loss.detach()) 

50 

51 

52def mean_global_loss( 

53 losses: Union[dict[str, torch.Tensor], torch.Tensor], 

54 current_token_counts: dict[str, torch.Tensor], 

55 device_mesh: "MeshContext", 

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

57 """Calculate global loss numerators using explicit mesh information. 

58 

59 FSDP divides gradients over its flattened DP+CP domain, so each local loss 

60 is weighted by valid tokens and multiplied by ``dp_size * cp_size``. 

61 Trainer divides accumulated gradients and reported losses by the actual 

62 optimizer-step token count afterwards. Empty shards must return finite losses. 

63 

64 Args: 

65 losses: A loss tensor or mapping of named loss tensors. 

66 current_token_counts: Token counts for the current micro batch. 

67 device_mesh: Trainer mesh context containing parallel sizes and groups. 

68 

69 Returns: 

70 Token-weighted loss numerators keyed by loss name. 

71 

72 Raises: 

73 ValueError: If sequence parallelism is enabled without a TP device mesh. 

74 """ 

75 loss_dict = {} 

76 dp_cp_mesh = device_mesh.dp_cp_mesh 

77 dp_cp_group = dp_cp_mesh.get_group() if dp_cp_mesh is not None else None 

78 sequence_parallel = device_mesh.sequence_parallel 

79 sequence_parallel_group = None 

80 sequence_parallel_size = 1 

81 if sequence_parallel: 

82 if device_mesh.device_mesh is None or "tp" not in device_mesh.device_mesh.mesh_dim_names: 

83 raise ValueError("Sequence parallelism requires a DeviceMesh with a 'tp' dimension.") 

84 sequence_parallel_group = device_mesh.device_mesh["tp"].get_group() 

85 sequence_parallel_size = dist.get_world_size(group=sequence_parallel_group) 

86 

87 if isinstance(losses, torch.Tensor): # text loss only 

88 losses = {"foundation_loss": losses} 

89 

90 if len({key.split("_loss", maxsplit=1)[0] for key in losses}) != 1: 

91 raise ValueError("Step gradient normalization requires losses sharing one token-count family") 

92 

93 for key, cur_loss in losses.items(): 

94 loss_name = key.split("_loss", maxsplit=1)[0] # foundation/image_decoder/** 

95 

96 cur_token_len = current_token_counts[f"{loss_name}_tokens"] 

97 if sequence_parallel: 

98 if _is_tracing_or_compiling(): 

99 # Graph tracing: ``.item()`` is data-dependent and unsupported 

100 # by the static-shape tracer; aggregate on tensors instead. 

101 cur_token_len = differentiable_all_reduce( 

102 cur_token_len, "sum", sequence_parallel_group 

103 ) 

104 else: 

105 cur_token_len = all_reduce(cur_token_len.item(), op="sum", group=sequence_parallel_group) 

106 

107 if _is_tracing_or_compiling(): 

108 # Graph-traced variant of the aggregation below: identical math 

109 # expressed with tensor collectives so the static-shape tracer can 

110 # capture it (no ``.item()``, no data-dependent zero-token guard). 

111 # ``differentiable_all_reduce`` dispatches to functional collectives 

112 # under tracing. BaseTrainer passes token-count tensors as explicit 

113 # graph inputs so packed/ignored-token counts remain live. 

114 local_weighted_loss = cur_loss.float() * cur_token_len 

115 backward_loss = local_weighted_loss * device_mesh.dp_size * device_mesh.cp_size 

116 global_weighted_loss = local_weighted_loss.detach() 

117 if dp_cp_mesh is not None and dp_cp_mesh.size() > 1: 

118 global_weighted_loss = differentiable_all_reduce( 

119 global_weighted_loss, "sum", dp_cp_group 

120 ) 

121 cur_loss = global_weighted_loss + (backward_loss - backward_loss.detach()) 

122 if sequence_parallel: 

123 cur_loss = cur_loss / sequence_parallel_size 

124 loss_dict[key] = cur_loss 

125 continue 

126 

127 cur_loss = _reduce_weighted_loss( 

128 cur_loss, 

129 cur_token_len, 

130 device_mesh, 

131 dp_cp_group, 

132 ) 

133 

134 if sequence_parallel: 

135 cur_loss = cur_loss / sequence_parallel_size 

136 

137 loss_dict[key] = cur_loss 

138 

139 return loss_dict