def get_lr(self):
        """Learning rate decay functions from:
              https://openreview.net/pdf?id=BJYwwY9ll pg. 4"""

        # Warmup.
        if self.warmup_iter > 0:
            if self.num_iters == self.warmup_iter and self.lr_decay_tokens is not None:
                self.warmup_tokens = self.num_tokens
            if self.num_iters <= self.warmup_iter:
                return float(self.start_lr) * self.num_iters / self.warmup_iter

        if self.lr_decay_tokens is None:
            # For any iterations larger than `self.end_iter`, use `self.min_lr`.
            if self.num_iters > self.end_iter:
                return self.min_lr
            # If we are done with the warmup period, use the decay style.
            current_iter = self.num_iters - self.warmup_iter
            decay_iter = self.end_iter - self.warmup_iter
            decay_ratio = float(current_iter) / float(decay_iter)
        else:
            if self.num_tokens > self.lr_decay_tokens:
                return self.min_lr
            current_tokens = self.num_tokens - self.warmup_tokens
            decay_tokens = self.lr_decay_tokens - self.warmup_tokens
            decay_ratio = float(current_tokens) / float(decay_tokens)
        assert decay_ratio >= 0.0
        assert decay_ratio <= 1.0

        if self.decay_style == 'linear':
            lr = self.start_lr * (1.0 - decay_ratio)
        elif self.decay_style == 'cosine':
            lr = self.start_lr / 2.0 * (math.cos(
                math.pi * decay_ratio) + 1)
        elif self.decay_style == 'exponential':
            # exp(-0.693) = 1/2
            lr = self.start_lr * math.exp(-0.693 * decay_ratio)
        else:
            lr = self.start_lr
        return max(lr, self.min_lr)
Learning rate decay functions from:
              https://openreview.net/pdf?id=BJYwwY9ll pg. 4
megatron_lm/megatron/learning_rates.py