Deepspeed Doesn't Help Save Vram

The first issue is that the minGPT example int the Model Parallel GPU Training seems to be problem once again, at least for stage 3 training.

However, when following the example for deepspeed stage 3 with offloading(on williamFalcon's fork of Karparthy's mingpt), get end up with this:

trainer = Trainer(gpus=1, precision=16, max_epochs=500,
                  gradient_clip_val=1.0, 
                  progress_bar_refresh_rate=1,
                  strategy=DeepSpeedPlugin(
                  stage=3,
                  offload_optimizer=True,  # Enable CPU Offloading
                  cpu_checkpointing=True,  # (Optional) offload activations to CPU
                  allgather_bucket_size=1e8, 
                  reduce_bucket_size=1e8
                ),)

Note: I had to remove the callback given that it was causing issues.
Given that my Vram is quite limited(GTX 1060 3GB), I tuned down both bucket sizes, as suggested, in half. Then I turned down the bucket size by an order of magnitude. Without deepspeed, I could have around a batch size of 48. However, even with deepspeed, I can still train with a batch size of 48.

Does deepspeed just not work on such low vram GPUs, or is pytorch lightning not instantiating deepspeed correctly?

cc @SeanNaren @awaelchli @rohitgr7

Start by reproducing the minGPT Model Parallel GPU Training example with the shown Trainer and DeepSpeedPlugin stage 3 configuration on a GTX 1060 3GB. Compare the batch size and VRAM behavior with and without offloading and reduced bucket sizes; done means determining whether the configuration is instantiated correctly or documenting the GPU limitation.