Sharded Training not saving memory
🐛 Bug

I am training an ImageGPT model, but am not seeing any less GPU memory being used when training with ddp sharded vs. without.
Environment info:
PTL: v1.1.8
PTL Bolts: v0.3.0
Pytorch: v1.7.1
Python: v3.7.7
Cuda: 10.2
Single AWS p3.8xlarge instance- it contains 4 Tesla V100 GPUs.
The code is very simple:
import pytorch_lightning as pl
from pytorch_lightning.plugins.sharded_plugin import DDPShardedPlugin
from pl_bolts.models import ImageGPT
from pl_bolts.datamodules import MNISTDataModule

  dm = MNISTDataModule('.', batch_size=32)
  # 230K parameters.
  model = ImageGPT(embed_dim=32, layers=16, heads=4, vocab_size=32, num_pixels=28)
  trainer = pl.Trainer(accelerator="ddp", gpus=4,) #plugins=[DDPShardedPlugin()])
  trainer.fit(model, dm)

When I remove the DDPShardedPlugin I am seeing GPU memory usage of ~13MiB

When I include the plugin, I am still seeing GPU memory usage exactly the same

I would expect the per-device memory usage with the plugin to be less.
@SeanNaren any idea on what's going on here? Thanks a lot for the help.