Not see desirable GPU memory saving when running DeepSpeedExamples/BingBertSquad

[wangli1426]

Hi ALL,

I run DeepSpeedExamples/BingBertSquad using docker image deepspeed/deepspeed:latest . I launched 16 ranks on two GPU servers, each of which is equipped with 8 V100 GPUs.

However, only 20% GPU memory saving, aka. 2.5GB per GPU vs 1.9 GB per GPU, was observed when I turned on zero_optimization and activation_checkpointing. Is this normal?

The config is attached as below.

{
  "train_batch_size": 256,
  "train_micro_batch_size_per_gpu": 16,
  "steps_per_print": 32,
  "gradient_accumulation_steps": 1,
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 3e-5,
      "weight_decay": 0.0,
      "bias_correction": false
    }
  },
  "gradient_clipping": 1.0,
  "fp16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 2
  },
  "activation_checkpointing": {
    "partition_activations": true,
    "cpu_checkpointing": false,
    "contiguous_memory_optimization": false,
    "number_checkpoints": null,
    "synchronize_checkpoint_boundary": false,
    "profile": true
    }
}

The network I used is Bert-large with 1024 hidden_size and 24 num_hidden_layers.

Any suggestion or feedback is highly appreciated.

Thank you.

[wangli1426]

Any suggestions?

[samyam]

@wangli1426 Thank you for trying out DeepSpeed.

For the Bert Large the max memory saving from ZeRO Stage 2 should be about 4GB.

Activation Checkpointing is not implemented in BingBertSquad so that will not make any difference.

How are you measuring the memory savings? nvidia-smi does not give a good indication of memory usage in case you are using that because torch caches a bunch of memory. You can print out the actual memory usage using torch.cuda.memory_allocated().

[wangli1426]

@samyam
Thank you for your response.

Yes, we noticed Activation Checkpointing is not implemented in BertBertSquad. We will try to implement it and will test what happens when Activation Checkpointing is enabled.

Yes, I measured the memory savings by Nvidia-smi. I will use torch.cuda.memory_allocated() to print the memory usage instead.

Thank you.