Make new optimizer more extensible, easier to integrate downstream for FSDP

[Zach Mueller (muellerzr)]

Description
This PR makes it easier for users to use FSDP with MS-AMP from their existing optimizers. This is especially beneficial for library authors, as currently we need to go through quite a bit to get the FSDP version of these optimizers working when a user passes in optim.Adam.

Instead we delegate the FSDPAdamW to an OptimWrapper, which calls an underlying optimizer as a passthrough. This lets us add in any logic that should be done before/after said logic easier, and it takes in a constructed Optimizer rather than being inherited.

Let me know what we think about this, currently I'm going through integrating FSDP and DeepSpeed w/ MS-AMP into Accelerate and found this to be a critical painpoint, as our users pass in normal PyTorch optimizers and don't create special versions themselves.

Yuxiang Yang (@tocean) Kan Wu (@wkcn) let me know what you two think :)

New working FSDP:

model, optimizer = ...
model, optimizer = msamp.initialize(model, optimizer, use_fsdp=True, weight_qtype=Dtypes.kfloat8_e4m3)
model = FP8FullyShardedDataParallel(model, use_orig_params=True, auto_wrap_policy=my_auto_wrap_policy)
optimizer = FSDPAdamW(optimizer)

[Zach Mueller (muellerzr)]

@microsoft-github-policy-service agree company="Hugging Face"

[Zach Mueller (muellerzr)]

Sorry for the extraneous pushes while I was figuring something out. Good to go now :)

[Zach Mueller (muellerzr)]

You can see our new accelerate benchmarking scripts here: https://github.com/huggingface/accelerate/tree/muellerzr-msamp-ds-fsdp/benchmarks/fp8/ms_amp

[Zach Mueller (muellerzr)]

Yuxiang Yang (@tocean) Kan Wu (@wkcn) any particular issues with this? :)

(Ideally it'd be great to include this in the next accelerate release on the 1st :) )

[Kan Wu (wkcn)]

Zach Mueller (@muellerzr) Thanks for your contribution!

The PR looks good to me.
Sorry that I am not at Microsoft and do not have the authorization to review and merge the pull request.

[Zach Mueller (muellerzr)]

Ack okay, I suppose we'll have to wait for Yuxiang Yang (@tocean) /Yifan Xiong (@abuccts) /guoshzhao to take a look. Thanks for the flag 🤗