[BREAKING][misc] feat: Abstract optimizer

[EduardDurech]

Abstract optimizer so can be used with whatever module and method a user wants, should be backwards compatible as default is torch.optim.AdamW, adds {actor_rollout_ref.actor,critic}.optim.{optimizer,optimizer_impl,override_optimizer_config}

# Default
optimizer_impl: torch.optim
optimizer: AdamW

# Example
optimizer_impl: torchao.optim
optimizer: _AdamW
override_optimizer_config:
  bf16_stochastic_round: true

Important: fsdp_sft_trainer optim aligned with FSDP optim optim.warmup_steps_ratio->optim.lr_warmup_steps_ratio

[gemini-code-assist]

Code Review

This pull request introduces a flexible optimizer abstraction, allowing users to specify any optimizer via configuration. This is a great enhancement for modularity. My review focuses on the implementation of the new build_optimizer function. I've identified a critical issue where the argument handling for the dynamically loaded optimizer is not robust and can lead to runtime TypeError exceptions. My suggestion involves using Python's inspect module to build the arguments dictionary safely, ensuring only valid parameters are passed to the optimizer's constructor. This will make the implementation more generic and prevent unexpected crashes.

[wuxibin89]

@EduardDurech Please resolve conflicts with main branch.

[EduardDurech]

@wuxibin89 I overwrote the PR #3692 as extra parameters should now be defined in {actor_rollout_ref.actor,critic}.optim.override_optimizer_config I'll comment on the PR

ci is unrelated to PR, tests passed #cf4cc6a6c60b2a21b1765825b83158ae6bea101b

cpu_unit_tests

huggingface_hub.errors.LocalEntryNotFoundError: An error happened while trying to locate the file on the Hub and we cannot find the requested files in the local cache. Please check your connection and try again or make sure your Internet connection is on.

sgl

RuntimeError: Maybe you called ray.init twice by accident? This error can be suppressed by passing in 'ignore_reinit_error=True' or by calling 'ray.shutdown()' prior to 'ray.init()'.

e2e_ascend

File "/usr/local/python3.10.17/lib/python3.10/site-packages/requests/models.py", line 978, in json
(TaskRunner pid=213242)     raise RequestsJSONDecodeError(e.msg, e.doc, e.pos)
(TaskRunner pid=213242) requests.exceptions.JSONDecodeError: Expecting value: line 1 column 1 (char 0)