Add capturable single-tensor RAdam

[mlazos]

🚀 The feature, motivation and pitch

Currently RAdam doesn't have a capturable implementation for cudagraphs. This is inconsistent with the other optimizer implementations. I implemented a foreach capturable implementation in #117912. It would be useful to extend this to single tensor as a starting point.

Alternatives

No response

Additional context

No response

cc @vincentqb @jbschlosser @albanD @janeyx99 @crcrpar @ezyang @chauhang @penguinwu

[MarouaneMaatouk]

Hi @mlazos @janeyx99, could you please share any docs or guidance to help understand what needs to be done to resolve this issue.
Is the task to switch to supported operations by cudagraphs when capture enabled, if that's the case where could I find them?
Maybe https://github.com/pytorch/pytorch/blob/main/torch/_dynamo/variables/builtin.py#L101?

[mlazos]

Hi @mlazos @janeyx99, could you please share any docs or guidance to help understand what needs to be done to resolve this issue. Is the task to switch to supported operations by cudagraphs when capture enabled, if that's the case where could I find them? Maybe main/torch/_dynamo/variables/builtin.py#L101?

Hi @MarouaneMaatouk,
The main guidance is to only do math using torch.* ops on tensor values (even for scalars) vs. a lot of math which is done with python int and float values. I recommend checking https://github.com/pytorch/pytorch/blob/1e78dc95a4aa270d8b70396c1476593b4cf5cf4d/torch/optim/adam.py as a good reference. Also feel free to ping with questions and put me as a reviewer on the PR. At a high level, you should wrap the python values with tensors.

[MarouaneMaatouk]

@mlazos I opened a PR: #118697