def step(self, increment, consumed_tokens):
        """Set lr for all parameters groups."""
        self.consumed_tokens = consumed_tokens
        self.num_steps += increment
        new_lr = self.get_lr()
        for group in self.optimizer.param_groups:
            group['lr'] = new_lr
        self.cur_lr = new_lr
Set lr for all parameters groups.
training/data_efficiency/gpt_finetuning/learning_rates.py