Can't reproduce 10B parameters GPT-2 training, CUDA out of memory

[ollmer]

Hi!
I'm trying to run 10B params model training on a 32 V100 GPUs, following the tutorial https://www.deepspeed.ai/tutorials/zero/#training-a-10b-parameter-gpt-2-model
But i consistently get CUDA out of memory even when i've reduced model size to 6B params and batch size to 1.
Here are my configs:

       --model-parallel-size 1 \
       --num-layers 40 \
       --hidden-size 3072 \
       --num-attention-heads 24 \
       --batch-size 1 \
       --seq-length 1024 \
       --max-position-embeddings 2048 \

deepspeed config

{
  "train_batch_size": 32,
  "gradient_accumulation_steps": 1,
  "steps_per_print": 100,  "gradient_clipping": 1.0,
  "wall_clock_breakdown": true,  
  "fp16": {    "enabled": true,    "loss_scale": 0,
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "min_loss_scale": 0.0000001
  },
  "zero_optimization": {
    "stage":2,
    "contiguous_gradients": true,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 50000000,
    "allgather_bucket_size": 500000000
  }
}

Hardware setup is 2 DGX-2.
Do i need to turn on activations checkpointing to make it work? Or is there some other caveats?

[tjruwase]

@ollmer Can you clarify how GPU count and RAM size?

[ollmer]

It runs on 32 V100 GPUs with 32Gb RAM each.

[ollmer]

Also worth mention that It's fp16 mode. So i expect optimizer state to require 120Gb of ram, which should be 4 Gb per GPU in zero stage 2 mode.

[tjruwase]

Can you please share your training log?

[tjruwase]

@ollmer Just wanted to follow up whether you could share your training log? Thanks.

[ollmer]

Sorry for delay. Here it is.
failed10b_train_2_cuda_oom.log

[tjruwase]

Thanks for sharing the log. Can you try the following?

1. Disable overlap_comm
2. Reduce allgather_bucket_size to 1e8 or 5e7

[ollmer]

I've tried it. Still got OOM.
failed10b_train_3_cuda_oom.log

[tjruwase]

Thanks for sharing the log. I will try to repro this issue on my side.

[samyam]

@ollmer the OOM you are seeing should be resolved by enabling activation_checkpointing in addition to enabling zero_offload.

[ollmer]

Hi @samyam!
I'm talking about zero-2 mode, not about offload to CPU. And that was exactly my question about checkpointing from the original post:

Do i need to turn on activations checkpointing to make it work? Or is there some other caveats?

I wonder what is the maximum size of the model that could be trained without checkpointing? Because it slows thing down a lot. The tutorial looks like it assumes that 10B model could be trained without checkpointing.

[samyam]

@ollmer I see now. Yes, on 32 V100 GPUs ZeRO-2 should work without offload, and yes you need to turn on activation checkpointing. Model size beyond a few billion parameters all need activation checkpointing to be enabled in addition to ZeRO-2 or ZeRO-Offload, which ever you might be using. We will clarify that on the tutorial.

The activation checkpointing will have about 33% overhead, often lower. Are you seeing a much larger overhead?