Merge pull request #678 from google:tflop_moe

PiperOrigin-RevId: 639194660
google · May 31, 2024 · f12ba54 · f12ba54
2 parents 3a441e8 + 288912b
commit f12ba54
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/MaxText/maxtext_utils.py b/MaxText/maxtext_utils.py
@@ -109,7 +109,7 @@ def calculate_tflops_training_per_device(config, log=True):
  if config.num_experts > 1:
  # MoE: brute force implementation
  gate_flops = 2 * config.per_device_batch_size * config.max_target_length * config.emb_dim * config.num_experts
- total_ffn_flops = gate_flops + config.num_experts * total_ffn_flops
+ total_ffn_flops = gate_flops + config.num_experts_per_tok * total_ffn_flops
 
  qkv_flops = (
  2