def step(self, step_num=None, token_num=None):
        """Set lr for all parameters groups."""
        args = get_args()
        if step_num is None:
            step_num = self.num_iters + 1
        if token_num is None:
            token_num = args.tokens
        self.num_iters = step_num
        self.num_tokens = token_num
        new_lr = self.get_lr()
        for group in self.optimizer.param_groups:
            group['lr'] = new_lr
Set lr for all parameters groups.
megatron_lm/megatron/learning_rates.py