Add support for DS comms

[Quentin-Anthony]

@awan-10 -- FYI

DeepSpeed no longer directly calls torch.distributed deepspeedai/DeepSpeed#1985. This commit updates Megatron-DeepSpeed to adhere to the new deepspeed.comm interface.

[Quentin-Anthony]

One thing I'm uncertain about is whether we should wrap calls in checks to whether args.deepspeed is set like in: https://github.com/microsoft/Megatron-DeepSpeed/blob/main/megatron/training.py#L114

E.g.

if args.deepspeed:
    deepspeed.comm.get_rank()
else:
    torch.distributed.get_rank()

However, this would require that args be propagated to all files where comms are used, which may get a bit messy. Thoughts?