allow specifying LR schedule in terms of tokens

[epwalsh]

This PR allows us to specify the LR schedule in terms of tokens instead of steps. For example, just change this:

scheduler:
  name: linear_with_warmup
  t_warmup: 5000
  t_max: 476837
  alpha_f: 0.1
  grad_clip_warmup_steps: 1000
  grad_clip_warmup_factor: 10.0

To this:

 scheduler:
   name: linear_with_warmup
+  units: tokens
-  t_warmup: 5000
-  t_max: 476837
+  t_warmup: 2e10
+  t_max: 2e12
   alpha_f: 0.1
-  grad_clip_warmup_steps: 1000
+  grad_clip_warmup_steps: 4e9
   grad_clip_warmup_factor: 10.0

The above two configurations are equivalent given a constant batch size after restarts, but the latter allows us to continue the same LR schedule while changing the batch size without any additional config changes.

This is backwards compatible, so you can make this change to your config and still restart from an older checkpoint.

[2015aroras]

Correct me if I'm wrong, but (maybe barring some weird edge cases) it seems that the conversion from steps to tokens (and from self.global_step to self.global_train_tokens_seen) is multiplying by self.tokens_per_batch. If that is the case, then it may be more readable if you make self.max_tokens return self.max_steps * self.tokens_per_batch, or some equivalent code.

[epwalsh]

I don't think that works if batch size has changed at some point.

[dirkgr]

Looks good. Did you test this in any way?

[epwalsh]

Looks good. Did you test this in any way?

@dirkgr, no, but I will before I merge.

[epwalsh]

Confirmed it's working as expected after a restart with 2x batch size. https://wandb.ai/ai2-llm/olmo-small-test?workspace=user-epwalsh