Corrects FLOPs formula as per 1093 (#1094)

* Update logging.py * Update NeoXArgs docs automatically --------- Co-authored-by: github-actions <[email protected]>
EleutherAI · Dec 6, 2023 · 050f560 · 050f560 · elsydony · Dec 17, 2023
1 parent a2b2020
commit 050f560
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/configs/neox_arguments.md b/configs/neox_arguments.md
@@ -111,7 +111,7 @@ Logging Arguments
 
 - **git_hash**: str
 
- Default = efaee8d
+ Default = bb1b145
 
  current git hash of repository
 

diff --git a/megatron/logging.py b/megatron/logging.py
@@ -92,17 +92,15 @@ def get_flops(neox_args, iter_time_s) -> float:
  hidden_size = neox_args.hidden_size
  num_layers = neox_args.num_layers
  ckpt_activations_factor = 4 if neox_args.checkpoint_activations else 3
- flops_calc1 = (
+ flops_per_iteration = (
  24
  * ckpt_activations_factor
  * batch_size
  * seq_len
  * num_layers
  * (hidden_size**2)
- * (1.0 + (seq_len / (6.0 * hidden_size)))
+ * (1.0 + (seq_len / (6.0 * hidden_size)) + (vocab_size / (16.0 * num_layers * hidden_size)))
  )
- flops_calc2 = vocab_size / (16.0 * num_layers * hidden_size)
- flops_per_iteration = flops_calc1 + flops_calc2
  return flops_per_iteration / (iter_time_s * world_size)