Refactoring FSDP.

[AdamLouly]

We are refactoring the FSDP implementation.

Fortunately, there aren't many changes required, as most of the modifications are in the Transformers Trainer. We don't need to duplicate those changes in the Optimum Trainer.

[JingyaHuang]

Hi @AdamLouly, thanks for the PR 🙏 .

The changes look good to me. The CI for ORT Training failed among which the nightly_test_trainer.py passed which means that the the ORTTrainer API is working, whereas nightly_test_examples.py , probably some training examples need to be updated as well.

Would you mind updating these training examples as well, otherwise we can get this PR in and target a new PR for updating those training examples.

cc. @prathikr

[JingyaHuang]

I think we can get this fix in and program another PR for updating those ORT training examples. Hi @AdamLouly, could you rebase the branch, as other maintainers have fixed some CIs which failed.

[JingyaHuang]

With the commits I submitted, examples are running well in the docker container I built locally, but the CI failed with following error:

Error: 2-26 13:55:32,610] torch.distributed.elastic.multiprocessing.api: [ERROR] failed (exitcode: -11) local_rank: 0 (pid: 280) of binary: /home/onnxruntimedev/miniconda3/bin/python

I will investigate further, but it's irrelevant to the fix so I will merge it anyway.

Thanks again for the contribution @AdamLouly !

[JingyaHuang]

LGTM