feat. Add more optimizer choice

[ZiyiTsang]

This pull request expands optimizer support in the codebase, allowing users to select between Adam, SGD, and a new AnyPrecision AdamW optimizer (adam_bf16). It also introduces a custom implementation of AnyPrecisionAdamW for improved mixed-precision training and updates documentation and CLI argument validation accordingly.

Optimizer Support Enhancements:

• Added support for sgd and adam_bf16 optimizer types in OptimizerConfig, and updated CLI argument validation to reflect these new options.
• Updated the base HuggingFace engine (areal/engine/base_hf_engine.py) to allow selection between AdamW, SGD, and the new AnyPrecisionAdamW optimizer, invoking the appropriate optimizer based on user configuration.
• Updated the Megatron engine (areal/experimental/megatron_engine.py) to allow both AdamW and SGD optimizers, passing the selected type to the optimizer configuration.

New Optimizer Implementation:

• Added a new AnyPrecisionAdamW optimizer in areal/utils/fsdp/optimizer.py, supporting flexible precision (bfloat16/float32) and optional Kahan summation for improved numerical stability during mixed-precision training.

AdamW_bf16/SGD optimizer maybe useful to avoid OOM, but less stability

Note
AnyPrecisionAdamW should not be changed as it is from other repo's implementation.
PR #366 should put after this PR

Optimizer Support

Optimizer	FSDP	Megatron
AdamW	✅	✅
SGD	✅	✅
AdamW_bf16	✅	❌

[ZiyiTsang]

/gemini review

[ZiyiTsang]

PR #366 should put after this PR

[ZiyiTsang]

Clear for review.

[garrett4wade]

LGTM, waiting for other reviewer's comments.

[nuzant]

LGTM, only one small issue.

[Copilot]

Pull Request Overview

This pull request expands optimizer support in the codebase, allowing users to select between Adam, SGD, and a new AnyPrecision AdamW optimizer (adam_bf16) for improved mixed-precision training.

• Adds support for sgd and adam_bf16 optimizer types across CLI validation, FSDP, and Megatron engines
• Implements a custom AnyPrecisionAdamW optimizer with flexible precision and optional Kahan summation
• Refactors optimizer creation from base class to engine-specific implementations

Reviewed Changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
docs/cli_reference.md	Updates documentation to reflect new optimizer choices and parameter applicability
areal/utils/fsdp/optimizer.py	Adds new AnyPrecisionAdamW optimizer implementation with precision control
areal/experimental/megatron_engine.py	Updates Megatron engine to support AdamW and SGD optimizers
areal/engine/ppo/actor.py	Minor formatting changes (whitespace additions)
areal/engine/fsdp_engine.py	Implements optimizer creation with support for all three optimizer types
areal/engine/base_hf_engine.py	Refactors optimizer creation to abstract method pattern
areal/api/cli_args.py	Updates CLI argument validation and adds version check logic

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[rchardx]

/gemini review

[ZiyiTsang]

/gemini review

[ZiyiTsang]

Done. Review again please @garrett4wade .

[garrett4wade]

LGTM