Enable torch.compile with ZeRO (Experimental) (#4878)
This PR enables `torch.compile` with ZeRO stages 1/2/3. You need to add
`compile` section in your DeepSpeed config. The fields in the section
are passed to `torch.compile`.
```json
"compile": {
"disable": false,
"backend": "inductor"
}
```
To enable a custom backend, you can pass the fully qualified name of the
backend function. For example, if you have a backend class `my_backend`
in `my_backend.py` in the current directory, you can enable it by
`"backend": "my_backend.my_backend"`. You can find an example in [a unit
test](https://github.com/microsoft/DeepSpeed/blob/eb9d4e06e9596f391aea305a6a5c6ec70cc28b58/tests/unit/runtime/compile/test_config.py#L116).
Currently we validated the results with Megatron-DeepSpeed. See the
[example](https://github.com/microsoft/Megatron-DeepSpeed/tree/tohtana/enable_compile/examples_deepspeed/compile)
for the details.
NOTICE: This PR is a draft. We will need to validate the coverage and
accuracy with many more examples.
---------
Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: Michael Wyatt <michaelwyatt@microsoft.com>
