Fix DeepSpeed (ZeRO2 + Pipeline Parallel)

[StellaAthena]

There is an issue with the DeepSpeed library that prevents you from using Pipeline Parallelism and ZeRO Stage 2 at the same time. @leogao2 has a rudimentary patch that allows the code to run (see here) but it causes a significant slowdown. We need to figure out how to do this better. For additional on the problem at hand, see #62

Profiling results from initial patch:

• patched, zero2+checkpoint+pipeline: samples/sec: 1159.741, max vram: 3245MiB
• patched, zero2+checkpoint: samples/sec: 1120.8568733324405, max vram: 1704MiB

[g-karthik]

@StellaAthena @leogao2 How is this a significant slowdown? Looks like the SamplesPerSec are quite high, w/ and w/o PP.

Also, what do you attribute the reduction in VRAM to in the case w/o PP?