torch.Tensor for optimizer parameters

[ad8e]

🚀 The feature, motivation and pitch

The LR allows torch.Tensor, see: #120934 (comment)

But the other parameters don't, such as AdamW's beta2.

This causes torch.compile to fail (fallback to eager) when I compile an optimizer with changing parameters, since torch.compile does not allow non-constant floats.

Allowing changing the other (non-LR) parameters would be helpful because:

1. beta1 (momentum) needs a scheduler as the eigenvalue distribution changes
2. beta2 needs a scheduler when finetuning a model, if the optimizer states aren't loaded or if the data distribution changes. Otherwise, a long warmup is required (~200-1000 steps depending on beta2). I'm currently having problems with DCP optimizer saving, so this matters to me.
3. weight decay needs a scheduler if the math is done correctly

Alternatives

Either torch.compile working with floats or AdamW allowing torch.tensor params would work for me.

Additional context

No response

cc @vincentqb @jbschlosser @albanD @janeyx99 @crcrpar