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

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"""Analytical model-FLOPs estimation for trainer throughput metrics.""" 

16 

17from __future__ import annotations 

18 

19from typing import Any, Optional, Tuple 

20 

21# Qwen3-MoE always projects its dense MLP and every routed expert through three 

22# GEMMs (gate, up, down). ``hidden_act`` only selects the elementwise activation 

23# applied to the gate branch, so it must not change the projection count. 

24_GATED_FFN_PROJECTIONS = 3.0 

25 

26 

27def _qwen3_moe_coefficients(config: Any) -> Optional[Tuple[float, float]]: 

28 """Split per-token FLOPs into a length-free part and a per-position part. 

29 

30 Causal attention charges ``query * seq_length`` per layer, so a sequence of 

31 length ``S`` costs ``tokens * base + S * S * attention`` FLOPs. Keeping the 

32 two coefficients apart lets callers accumulate ``sum(S_i ** 2)`` over the 

33 real attention boundaries instead of assuming one representative length. 

34 

35 Args: 

36 config: Hugging Face model configuration. 

37 

38 Returns: 

39 ``(base, attention)`` FLOPs per token and per token-position, or 

40 ``None`` when the model family is not supported. 

41 """ 

42 if getattr(config, "model_type", None) != "qwen3_moe": 

43 return None 

44 

45 hidden = int(config.hidden_size) 

46 layers = int(config.num_hidden_layers) 

47 heads = int(config.num_attention_heads) 

48 kv_heads = int(getattr(config, "num_key_value_heads", None) or heads) 

49 head_dim = int(getattr(config, "head_dim", None) or hidden // heads) 

50 intermediate = int(config.intermediate_size) 

51 moe_intermediate = int(config.moe_intermediate_size) 

52 topk = int(config.num_experts_per_tok) 

53 vocab = int(config.vocab_size) 

54 

55 query = heads * head_dim 

56 key_value = kv_heads * head_dim 

57 

58 # Per token, per standard-attention layer: QKV and output projections. 

59 attention_base = 6.0 * (hidden * (query + key_value + key_value) + query * hidden) 

60 # Causal attention only attends the lower triangle, i.e. ``query * S``. 

61 attention_position = 6.0 * query 

62 

63 projections = _GATED_FFN_PROJECTIONS 

64 dense_mlp = 6.0 * hidden * intermediate * projections 

65 moe_mlp = 6.0 * hidden * moe_intermediate * topk * projections 

66 

67 # Match Qwen3MoeDecoderLayer: MoE when the layer is not in 

68 # ``mlp_only_layers`` and ``(layer_idx + 1) % decoder_sparse_step == 0``. 

69 mlp_only = set(getattr(config, "mlp_only_layers", None) or ()) 

70 sparse_step = max(1, int(getattr(config, "decoder_sparse_step", 1) or 1)) 

71 num_moe = sum( 

72 1 

73 for layer_idx in range(layers) 

74 if layer_idx not in mlp_only 

75 and int(getattr(config, "num_experts", 0) or 0) > 0 

76 and (layer_idx + 1) % sparse_step == 0 

77 ) 

78 num_dense = layers - num_moe 

79 

80 logits = 6.0 * hidden * vocab 

81 base = ( 

82 attention_base * layers 

83 + dense_mlp * num_dense 

84 + moe_mlp * num_moe 

85 + logits 

86 ) 

87 return base, attention_position * layers 

88 

89 

90def estimate_flops_per_token(config: Any, seq_length: float) -> Optional[float]: 

91 """Estimate training FLOPs per token for a decoder-only model. 

92 

93 The estimate counts forward, weight-gradient and data-gradient GEMMs 

94 (3x) and two FLOPs per multiply-accumulate (2x), so every GEMM carries a 

95 factor of six. Causal attention counts only the lower triangle via 

96 ``q * seq_length``. 

97 

98 Args: 

99 config: Hugging Face model configuration. 

100 seq_length: Representative tokens per sequence. Use this helper only 

101 when every sequence in the step has that length; otherwise prefer 

102 :func:`estimate_step_flops`, which accumulates the attention term 

103 over the real sequence boundaries. 

104 

105 Returns: 

106 Estimated FLOPs per token, or ``None`` when the model family is not 

107 supported. Communication, recomputation, normalization, dropout, 

108 RoPE and optimizer work are not included. 

109 """ 

110 coefficients = _qwen3_moe_coefficients(config) 

111 if coefficients is None: 

112 return None 

113 base, attention_position = coefficients 

114 return base + attention_position * float(seq_length) 

115 

116 

117def estimate_step_flops( 

118 config: Any, 

119 input_tokens: float, 

120 sum_seq_len_sq: float, 

121) -> Optional[float]: 

122 """Estimate training FLOPs for one optimizer step. 

123 

124 Attention FLOPs scale with the square of each sequence length, so a step is 

125 described by the model-input token count ``sum(S_i)`` and the length squares 

126 ``sum(S_i ** 2)``. For packed batches, ``S_i`` are the packed segment 

127 lengths; for plain batches they are the per-row token counts. 

128 

129 Args: 

130 config: Hugging Face model configuration. 

131 input_tokens: Model-input tokens consumed by the step, i.e. tokens that 

132 reach attention and the FFN, including prompt tokens masked out of 

133 the loss. 

134 sum_seq_len_sq: Sum of the squared length of every sequence that 

135 reaches attention in the step. 

136 

137 Returns: 

138 Estimated step FLOPs, or ``None`` when the model family is not 

139 supported. Communication, recomputation, normalization, dropout, 

140 RoPE and optimizer work are not included. 

141 """ 

142 coefficients = _qwen3_moe_coefficients(config) 

143 if coefficients is None: 

144 return None 

145 base, attention_position = coefficients 

146 return base * float(input_tokens) + attention_position * float(sum_seq_len_sq)