Add support for pytorch 2.0 ?

[guozhiyao]

[Quentin-Anthony]

This likely won't be explored by us until DeepSpeed supports it. torch.compile makes overlapping compute with communication difficult for now since hooks are broken.

Contributions along this angle are welcome though!