[DeepSpeed] unable to increase batch size from 4 for T5-3b with 2x 32GB V100 GPUs
Hi,

I'm trying T5-3b with DeepSpeed on 2 V100-32GB GPU's. But I'm unable to increase batch size beyond 4 with max i/p sequence length of 512 and max o/p sequence length of 4. 
Previously I tried with t5.parallelize() [ i.e, without DeepSpeed ] on same setup and was able to train with batch size of 2. 

Below are my training args and DeepSpeed's config - 
```
training_args = Seq2SeqTrainingArguments(
    output_dir='./seq_out/results',
    overwrite_output_dir=True,
    evaluation_strategy="epoch",
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    learning_rate=3e-5,
    weight_decay=0.01,
    num_train_epochs=2,    
    warmup_steps=500,             
    logging_dir='./seq_out/logs',          
    logging_steps=10,
    load_best_model_at_end=True,
    deepspeed='ds_config.json'
)
```
```
{
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "hysteresis": 2,
        "min_loss_scale": 1
    },

   "zero_optimization": {
       "stage": 2,
       "allgather_partitions": true,
       "allgather_bucket_size": 1.5e8,
       "overlap_comm": true,
       "reduce_scatter": true,
       "reduce_bucket_size": 1.5e8,
       "contiguous_gradients": true,
       "cpu_offload": true
   },

   "zero_allow_untested_optimizer": true,

   "optimizer": {
     "type": "AdamW",
     "params": {
       "lr": 3e-5,
       "betas": [
         0.8,
         0.999
       ],
       "eps": 1e-8,
       "weight_decay": 3e-7
     }
   },

   "scheduler": {
     "type": "WarmupLR",
     "params": {
       "warmup_min_lr": 0,
       "warmup_max_lr": 3e-5,
       "warmup_num_steps": 500
     }
   },

    "steps_per_print": 2000,
    "wall_clock_breakdown": false
}
```

Observed similar behavior with T5-large as well . I was able to train with 14 batch size ( same i/p, o/p seq length and config.json as above ) on single GPU ( 32 GB V100 ) BUT when executing the same on 2x 32 GB GPUs, I was not able to go beyond 14 batch size( which I was able to train with 1 GPU itself) and memory from both the GPUs was consumed (31 GB and 28 GB).

Reducing `allgather_bucket_size ` and `reduce_bucket_size` didn't help in increasing batch size.

But I expected more batch size with DeepSpeed and CPU offloading. Is this fine or am I making something wrong which is hindering deepspeed's capability..?
