Severe performance degradation with deepspeed > v.0.3.12

[g-karthik]

Creating a new issue with more details around what I reported in the comments of an unrelated issue #1057

I have been training a 8B parameter GPT-2 model with DeepSpeed Stage 2 enabled (micro-batch size 2, world size 64).

I used this base image nvcr.io/nvidia/pytorch:21.03-py3 and installed the latest DeepSpeed on top of it. It has pytorch 1.9, python 3.8, NCCL 2.8.4.

I was seeing the following stats per iteration when profiling:

forward: 294.88 | backward: 17065.34 | backward_inner: 15988.85 | backward_allreduce: 1076.32 | step: 2102.37
SamplesPerSec: ~7.5

I found this to contradict my previous tests which used older dependencies and so I downgraded deepspeed to v0.3.12, and then got the following stats (keeping all else equal):

forward: 294.88 | backward: 3244.52 | backward_inner: 3051.20 | backward_allreduce: 193.21 | step: 1651.08
SamplesPerSec: ~33.4

I don't use a user-defined optimizer and I see from my logs that DeepSpeed's FusedAdam is picked in both cases.

Clearly, the latest deepspeed seems to be severely hurting performance, as seen by backward time (a difference of ~14000 milliseconds).

I want to upgrade to the latest deepspeed, what is the cause of this degradation in performance with the latest deepspeed?

@jeffra

EDIT: Just a correction, it seems from my logs that DeepSpeed's FusedAdam is picked in the second case, but DeepSpeedCPUAdam is picked in the first case. That may possibly explain the difference in backward times. But this is despite my usage of "type": "adam" for the optimizer in the DeepSpeed config JSON in both cases. Why does this inconsistency exist?

[jeffra]

Hi @g-karthik, can you share your entire ds_config.json? I am a bit surprised by this actually. Also out of curiosity is this using Megatron or huggingface GPT-2?

[g-karthik]

Hi @jeffra yes it's HF GPT-2. Here's the DS config I used in both cases:

{
  "train_batch_size": 640,
  "gradient_accumulation_steps": 5,
  "gradient_clipping": 1.0,
  "optimizer": {
    "type": "adam",
    "params": {
      "lr": 9.25e-5
    }
  },
  "fp16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu"
    },
    "overlap_comm": true,
    "allgather_bucket_size": 500000000
  },

  "sparse_gradients": false,

  "activation_checkpointing": {
    "partition_activations": false,
    "contiguous_memory_optimization": false,
    "number_checkpoints": 50,
    "cpu_checkpointing": false
  },

  "wall_clock_breakdown": true,

  "flops_profiler": {
    "enabled": true,
    "profile_step": 30,
    "module_depth": -1,
    "top_modules": 3,
    "detailed": true
  }
}

[jeffra]

Hi @g-karthik, thanks for this context. If I wanted to try and reproduce this what version/commit should I use from https://github.com/huggingface/transformers?

This seems really strange, I would expect your ds_config.json to use DeepSpeedCPUAdam and not FusedAdam.

[g-karthik]

@jeffra you should be able to reproduce this with modeling_gpt2.py from version 4.5.1.

If "type": "adam" is always supposed to point to DeepSpeedCPUAdam, then what type would point to FusedAdam? And seeing as the backward time reduces from ~17 seconds to ~3 seconds (from above), I am struggling to imagine why I would want to use DeepSpeedCPUAdam anyway. Perhaps by leveraging CPU Adam, I could fit a larger model size but what would the point of doing that be if it comes at such a significant cost to backward time?