[FEATURE] DeepSpeed/ZeRO Support #490
Description
Along with updated training/validation components in #458 for TPU support, support use of DeepSpeed/ZeRO
- https://pytorch.org/docs/master/distributed.optim.html#torch.distributed.optim.ZeroRedundancyOptimizer
- https://github.com/microsoft/DeepSpeed
It would be fairly easy to support w/ current training code, however the cahnges in #458 are going to be so significant that I do not want to continue two branches of code.