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

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"""Global loss metrics reduction across data/context parallel ranks.""" 

15 

16from __future__ import annotations 

17 

18__all__ = ["mean_global_loss"] 

19 

20from typing import TYPE_CHECKING, Union 

21 

22import torch 

23import torch.distributed as dist 

24 

25from hyper_parallel.core.utils.communication import ( 

26 _is_tracing_or_compiling, 

27 differentiable_all_reduce, 

28) 

29from hyper_parallel.trainer.runtime.distributed import all_reduce 

30 

31if TYPE_CHECKING: 

32 from hyper_parallel.distributed.mesh import MeshContext 

33 

34 

35def _reduce_weighted_loss(cur_loss, cur_token_len, all_reduced_len, loss_name, device_mesh, dp_cp_group): 

36 """Reduce a token-weighted loss while preserving the local backward graph.""" 

37 if all_reduced_len != 0: 

38 local_weighted_loss = cur_loss * cur_token_len 

39 backward_loss = local_weighted_loss / all_reduced_len * device_mesh.dp_size * device_mesh.cp_size 

40 global_weighted_loss = all_reduce( 

41 local_weighted_loss.detach().item(), 

42 op="sum", 

43 group=dp_cp_group, 

44 ) 

45 global_mean = cur_loss.new_tensor(global_weighted_loss / all_reduced_len) 

46 return backward_loss + global_mean - backward_loss.detach() 

47 

48 if not torch.allclose(cur_loss, torch.zeros_like(cur_loss)): 

49 raise ValueError(f"The all_reduced_len for {loss_name}_tokens is 0, but the cur_loss is not 0: {cur_loss}") 

50 return cur_loss 

51 

52 

53def mean_global_loss( 

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

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

56 step_token_counts: dict[str, torch.Tensor], 

57 device_mesh: "MeshContext", 

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

59 """Calculate the global mean loss using explicit mesh information. 

60 

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

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

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 step_token_counts: Token counts for the local optimizer step. 

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

69 

70 Returns: 

71 Token-weighted loss tensors keyed by loss name. 

72 

73 Raises: 

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

75 """ 

76 loss_dict = {} 

77 dp_cp_mesh = device_mesh.dp_cp_mesh 

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

79 sequence_parallel = device_mesh.sequence_parallel 

80 sequence_parallel_group = None 

81 sequence_parallel_size = 1 

82 if sequence_parallel: 

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

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

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

86 sequence_parallel_size = dist.get_world_size(group=sequence_parallel_group) 

87 

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

89 losses = {"foundation_loss": losses} 

90 

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

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

93 

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

95 if sequence_parallel: 

96 if _is_tracing_or_compiling(): 

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

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

99 cur_token_len = differentiable_all_reduce( 

100 cur_token_len, "sum", sequence_parallel_group 

101 ) 

102 else: 

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

104 

105 if _is_tracing_or_compiling(): 

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

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

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

109 # ``differentiable_all_reduce`` dispatches to functional collectives 

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

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

112 step_len = differentiable_all_reduce( 

113 step_token_counts[f"{loss_name}_tokens"], "sum", dp_cp_group 

114 ) 

115 local_weighted_loss = cur_loss * cur_token_len 

116 backward_loss = local_weighted_loss / step_len * device_mesh.dp_size * device_mesh.cp_size 

117 global_weighted_loss = differentiable_all_reduce( 

118 local_weighted_loss.detach(), "sum", dp_cp_group 

119 ) 

120 global_mean = global_weighted_loss / step_len 

121 cur_loss = backward_loss + global_mean - 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 all_reduced_len = all_reduce( 

128 step_token_counts[f"{loss_name}_tokens"].item(), 

129 op="sum", 

130 group=dp_cp_group, 

131 ) 

132 

133 cur_loss = _reduce_weighted_loss( 

134 cur_loss, 

135 cur_token_len, 

136 all_reduced_len, 

137 loss_name, 

138 device_mesh, 

139 dp_cp_group, 

140 ) 

141 

142 if sequence_parallel: 

143 cur_loss = cur_loss / sequence_parallel_size 

144 

145 loss_dict[key] = cur_loss 

146 

147 return loss_dict