Fixed AnnealingLR Class and Cosine Decay Schedule (#1008)

* Fixed AnnealingLR Class and Cosine Decay Schedule * Update NeoXArgs docs automatically --------- Co-authored-by: github-actions <[email protected]>
EleutherAI · Aug 7, 2023 · 009018e · 009018e
1 parent 8b3f1c2
commit 009018e
Show file tree

Hide file tree

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

diff --git a/megatron/learning_rates.py b/megatron/learning_rates.py
@@ -64,7 +64,7 @@ def get_lr(self):
  """Learning rate decay functions from:
  https://openreview.net/pdf?id=BJYwwY9ll pg. 4"""
 
- num_iters_ = min(self.num_iters, self.end_iter - self.warmup_iter)
+ num_iters_ = self.num_iters
  # Warmup.
  if self.warmup_iter > 0 and self.num_iters <= self.warmup_iter:
  return float(self.start_lr) * num_iters_ / self.warmup_iter
@@ -73,8 +73,8 @@ 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":
- lr = (
- self.start_lr
+ lr = self.min_lr + (
+ (self.start_lr-self.min_lr)
  / 2.0
  * (math.cos(math.pi * num_iters_ / self.end_iter) + 1)
  )