CPU Offload is harmful to training convergence?

[robinn37]

Hi,

I've tried Deepspeed these days and it seems CPU offload doesn't work properly. My experiments are based on 1GB corpus and training GPT-2 with 6 layers. When CPU offload is enabled, the training loss stuck at around 4.0. I tried to update the lr, batch size, fp16, sequene len, activation_checkpointing and torch adam, but nothing helped. However, after switched off CPU offload, everything went well. The training loss and eval loss can reached below 2.0.

Anyone can shed some lights on this? Any args or configurations should be updated together with CPU offload?

-------- running configurations with CPU offload enabled -------------
{
"train_batch_size":32,
"train_micro_batch_size_per_gpu": 32,
"steps_per_print": 1,
"zero_optimization": {
"stage": 2,
"cpu_offload": true,
"reduce_bucket_size": 50000000
},
"zero_allow_untested_optimizer": false,
"gradient_clipping": 1.0,
"tensorboard": {
"enabled": true,
"output_path": "xxxx",
"job_name": "xxxxx"
},
"fp16": {
"enabled": true,
"loss_scale": 4096,
"loss_scale_window": 1000,
"hysteresis": 2,
"min_loss_scale": 1
},
"activation_checkpointing": {
"partition_activations": true,
"contiguous_memory_optimization": true,
"cpu_checkpointing": true
},
"wall_clock_breakdown": true
}

----- args ----
--num-layers 6
--hidden-size 512
--num-attention-heads 4
--seq-length 128
--max-position-embeddings 1024
--train-iters 20000
--resume-dataloader
--train-data webtext
--lazy-loader
--tokenizer-type GPT2BPETokenizer
--split 949,50,1
--distributed-backend nccl
--lr 0.00005
--no-load-optim
--lr-decay-style cosine
--weight-decay 1e-2
--clip-grad 1.0
--warmup .01
--checkpoint-activations
--deepspeed-activation-checkpointing
--fp16
--log-interval 50
--vocab-size 50257
--eval-interval 100
--cpu-optimizer
--cpu_torch_adam

------ training scripts ------
Use the training scripts in Megatron-LM from DeepSpeedExamples.

[robinn37]

training loss snapshots:

eval loss snapshots:

[RezaYazdaniAminabadi]

Hi @robinn37

Thanks for trying CPU offload and sharing the issue.
Yes, you are right there is difference in the convergence and we are aware of this!
The issue is the difference in the optimizer algorithm between baseline (fused-Adam) and (DeepSpeed-CPUAdam)! If you look at the apex fusedAadam implementation (https://github.com/NVIDIA/apex/blob/master/apex/optimizers/fused_adam.py#L63), you will see that it is by default using the AdamW algorithm, however we were using Adam. We just merged a PR which adds this implementation for DeepSpeed-CPUAdam (https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/ops/adam/cpu_adam.py#L58) and solves the convergence issue for the GPT2 training.
Please check in the new changes and try this again.

Thanks,
Reza

[RezaYazdaniAminabadi]

Please also check in the new changes for the DeepSpeedExample branch as we have worked on changing the scripts and configuration files too.

Thanks
Reza

[robinn37]

@RezaYazdaniAminabadi Thanks for your quick reply. I tried the latest code with the adamW update in deepspeed. And also tried the configurations update in DeepSpeedExample. However, the result has no difference.

------- deepspeed configurations --------------
"train_micro_batch_size_per_gpu": 32,
"gradient_accumulation_steps": 5,
"steps_per_print": 100,
"optimizer": {
"type": "Adam",
"params": {
"lr": 0.00015,
"weight_decay": 1e-2
}
},
"zero_optimization": {
"stage": 2,
"cpu_offload": true,
"reduce_bucket_size": 50000000
},
"gradient_clipping": 1.0,
"fp16": {
"enabled": true,
"loss_scale": 0,
"loss_scale_window": 1000,
"hysteresis": 2,
"min_loss_scale": 1
},
"activation_checkpointing": {
"partition_activations": true,
"contiguous_memory_optimization": true,
"cpu_checkpointing": true
},
"wall_clock_breakdown": false,

-------- startup script args -----------
gpt_options="
--model-parallel-size ${MP_SIZE}
--num-layers 6
--hidden-size 512
--num-attention-heads 4
--batch-size 32
--seq-length 128
--max-position-embeddings 1024
--train-iters 40000
--resume-dataloader
--train-data webtext
--lazy-loader
--tokenizer-type GPT2BPETokenizer
--split 949,50,1
--distributed-backend nccl
--lr 0.00015
--no-load-optim
--lr-decay-style cosine
--weight-decay 1e-2
--clip-grad 1.0
--warmup .01
--checkpoint-activations
--deepspeed-activation-checkpointing
--fp16
--log-interval 5
--vocab-size 50257
--eval-interval 100
"

Here are the training & eval loss curves for your information.

• The grey line is training without cpu-offload
• The green line is enable cpu-offload and adam with weight decay
• The orange line is same as the green one, but remove --cpu-optimizer from the script

[robinn37]

Use torch.optim.AdamW directly works well.

[RezaYazdaniAminabadi]

This is pretty strange! I have tried the same experiments and saw the same result! I will try using the same configuration as you and I will see if I can repro the same result.
Can you share the log of your previous runs? Also, I wonder how you directly used torch.optim.adamW, did you enable that using the deepspeed json config?

Thanks,
Reza

[robinn37]

I update the get_optimizer() in Megatron-LM/pretrain_gpt2.py:

+    #optimizer = Adam(param_groups,
+    #                 lr=args.lr, weight_decay=args.weight_decay)
+    optimizer = torch.optim.AdamW(param_groups,
+                                  lr=args.lr, weight_decay=args.weight_decay)

[RezaYazdaniAminabadi]

Hi @robinn37

Thanks for the update.
Could you please tell me how many GPUs and what model-parallelism degree you are using here, so that I can reproduce the issue with using the DeepSpeed-CPUAdam?

Thanks,
Reza

[robinn37]

2 GPU, MP degree = 1

[KeyKy]

@RezaYazdaniAminabadi I also have this problem. When cpu_offload is true, it does not converge. deepspeedai/DeepSpeedExamples#153

@robinn37 did you solve it ?