[RL] Add fp16 training support

[krammnic]

As discussed with @felipemello1 adding the support of the mixed precision training.

Parameters and optimizer state that are preserved in fp32, while FSDP casts parameters in fp16. For the PR the support is currently limited with Qwen model. Furthermore, the additional changes will be required to support MoE. Several new tests were introduced for the proposed changes.

From the non-trivial parts, I believe the only point which is worth clarification for this PR is a correction that is introduced.

The standard flow which I believed to be correct was quite simple:

loss.backward() -> $\nabla$ G
scaler.unscale_() -> $\nabla$ G / S

Therefore, it was important to introduce the correction:

(loss * S).backward -> $\nabla$ S * G
scaler.unscale_() -> $\nabla$ G

In the second case clipping will be applied to the unscaled gradients and optimization step will be correct!

Why it is still draft?

I haven't conducted the E2E run to verify that everything works correctly.

[pytorch-bot]

The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:

• ciflow/8gpu
• ciflow/rl

Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows.

[pytorch-bot]

The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:

• ciflow/8gpu
• ciflow/rl

Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows.