Learning rate decay functions from: https://openreview.net/pdf?id=BJYwwY9ll pg. 4
(self)
| 68 | print_rank_0("> learning rate decay style: {}".format(self.decay_style)) |
| 69 | |
| 70 | def get_lr(self): |
| 71 | """Learning rate decay functions from: |
| 72 | https://openreview.net/pdf?id=BJYwwY9ll pg. 4""" |
| 73 | |
| 74 | # Use linear warmup for the initial part. |
| 75 | if self.warmup_steps > 0 and self.num_steps <= self.warmup_steps: |
| 76 | if self.num_steps == self.warmup_steps and self.decay_tokens is not None: |
| 77 | self.warmup_tokens = self.num_tokens |
| 78 | return self.max_lr * float(self.num_steps) / float(self.warmup_steps) |
| 79 | |
| 80 | # If the learning rate is constant, just return the initial value. |
| 81 | if self.decay_style == "constant": |
| 82 | return self.max_lr |
| 83 | |
| 84 | if self.decay_tokens is None: |
| 85 | # step-based decay |
| 86 | |
| 87 | # For any steps larger than `self.decay_steps`, use `self.min_lr`. |
| 88 | if self.num_steps > self.decay_steps: |
| 89 | return self.min_lr |
| 90 | |
| 91 | # If we are done with the warmup period, use the decay style. |
| 92 | num_steps_ = self.num_steps - self.warmup_steps |
| 93 | decay_steps_ = self.decay_steps - self.warmup_steps |
| 94 | decay_ratio = float(num_steps_) / float(decay_steps_) |
| 95 | else: |
| 96 | # token-based decay |
| 97 | |
| 98 | if self.num_tokens > self.decay_tokens: |
| 99 | return self.min_lr |
| 100 | num_tokens_ = self.num_tokens - self.warmup_tokens |
| 101 | decay_tokens_ = self.decay_tokens - self.warmup_tokens |
| 102 | decay_ratio = float(num_tokens_) / float(decay_tokens_) |
| 103 | assert decay_ratio >= 0.0 |
| 104 | assert decay_ratio <= 1.0 |
| 105 | delta_lr = self.max_lr - self.min_lr |
| 106 | |
| 107 | if self.decay_style == "linear": |
| 108 | coeff = 1.0 - decay_ratio |
| 109 | elif self.decay_style == "cosine": |
| 110 | coeff = 0.5 * (math.cos(math.pi * decay_ratio) + 1.0) |
| 111 | else: |
| 112 | raise Exception("{} decay style is not supported.".format(self.decay_style)) |
| 113 | |
| 114 | return self.min_lr + coeff * delta_lr |
| 115 | |
| 116 | def step(self, increment, token_num=None): |
| 117 | """Set lr for all parameters groups.""" |