Fixed final value of cosine decay lr (#1011)

* Fixed final value of cosine decay lr * Update NeoXArgs docs automatically --------- Co-authored-by: github-actions <[email protected]>
EleutherAI · Aug 12, 2023 · eda5aae · eda5aae
1 parent 009018e
commit eda5aae
Show file tree

Hide file tree

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

diff --git a/megatron/learning_rates.py b/megatron/learning_rates.py
@@ -73,10 +73,11 @@ def get_lr(self):
  if self.decay_style == "linear":
  lr = self.start_lr * (self.end_iter - num_iters_) / self.end_iter
  elif self.decay_style == "cosine":
+ end_iter_ = self.end_iter - self.warmup_iter
  lr = self.min_lr + (
  (self.start_lr-self.min_lr)
  / 2.0
- * (math.cos(math.pi * num_iters_ / self.end_iter) + 1)
+ * (math.cos(math.pi * num_iters_ / end_iter_) + 1)
  )
  elif self.decay_style == "exponential":
  # exp(-0.693) = 1/2