Patch ORTTrainer's compatibility with DeepSpeed

[JingyaHuang]

What does this PR do?

Fixes #145 and #146

Contents

• Add missing dependencies
• Fix compatibility with deepspeed
• Fix compatibility with fairscale(simple ✅ dp2/dp3❌)
• Test ZeRO: stage 1
• Test ZeRO: stage 2
• Test ZeRO: stage 3 ❌
• Test ZeRO: NVMe support ❌
• Test ZeRO: BF16 ❌
• ZeRO Inference when ort inference enabled(need to coordinate with transformers.onnx.export_pytorch)

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[jambayk]

Hi, I saw that this PR is linked to issue #145 I opened.

I do not think this resolves the issue which came from model being used as if it's of type ORTModule. But it isn't when deepspeed or any other wrappers such as DDP are used.

EDIT: I apologize if this is still work in progress and you had a fix in the works. Please ignore this comment if so. Just wanted to make sure I provided enough context for the issue.

[JingyaHuang]

Hi, I saw that this PR is linked to issue #145 I opened.

I do not think this resolves the issue which came from model being used as if it's of type ORTModule. But it isn't when deepspeed or any other wrappers such as DDP are used.

EDIT: I apologize if this is still work in progress and you had a fix in the works. Please ignore this comment if so. Just wanted to make sure I provided enough context for the issue.

Hi @jambayk , no worries, I opened this PR so that there will be more transparency on the progress. Sorry for the confusion that it might have caused, and thanks a lot for the information on the issue.

[JingyaHuang]

The compatibility of ORTTrainer and DeepSpeed/Fairscale is implemented. Check the result of the tests here.

[lewtun]

Thanks for adding these warnings @JingyaHuang ! LGTM 🚀 !

[lewtun]

I wonder if we should check the transformers version and then raise a warning if the detected version doesn't match the required one for ONNX export on CUDA?

For now, this warning is OK but we'll probably want to revisit this once we bump transformers with your PR :)

[JingyaHuang]

@lewtun Yes, exactly. I put it this way since I am not sure in which version of transformers will include it. Will check the transformers version once I have the information.

[lewtun]

Sounds good!

[JingyaHuang]

PR#17183 to support ONNX export on CUDA. Will refactor the code once transformers includes it in a release.

[JingyaHuang]

Merge this to enable DeepSpeed with ORTTrainer, and pay attention to the following two points:

• If need to export mixed-precision trained(fp16) ONNX model, the export need to be done on CUDA. And this demand transformers version > 4.19.2(will be included in the next release).
• DeepSpeed BF16 is not compatible with onnxruntime-training. Will track the progress from ONNX Runtime side on supporting bf16 for essential ops.