fix mfu

PrimeIntellect-ai · Sep 22, 2024 · fc603ef · fc603ef
1 parent 293dfbd
commit fc603ef
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/zeroband/train.py b/src/zeroband/train.py
@@ -207,7 +207,7 @@ def train(config: Config):
             time_taken = time.time() - beginning_step_time
             tokens_per_second = config.data.seq_length * config.optim.batch_size / time_taken
 
-            mfu = 100 * num_flop_per_token * tokens_per_second / gpu_peak_flops
+            mfu = 100 * num_flop_per_token * tokens_per_second / gpu_peak_flops / world_info.local_world_size
 
             metrics = {
                 "Loss": loss_batch.item(),