deepspeedai / deepspeedai/DeepSpeed
Still OOM when enabling deepspeed and error during backward
@mrwyattii is already working on this.
Since Apr 21, 2023.
- Dominant language
- Python
- Stars
- 43.1k
- Forks
- 5k
- Avg merge
- 4d 15h
- Merged PRs (30d)
- 112
Description
I have several foundings and confusion about GPU memory during using the deepspeed. The environment I used is PyTorch-1.8.1, deepspeed-0.9.0:
(1) I am trying to convert a single-gpu cv-training project to support multi-gpu running via deepspeed, however I found that the deepspeed cmd seems to independently excute the project on two GPUs, because I found all the printed logs (such as iteration-loss info) are twice as single-gpu version, when using 4 gpus, all the logs occurred 4 times, Also the GPU memory consumption is almost same on each gpu. So how can I confirm that the model is trained in distributed way instead of independently?
(2) I have enabled Zero-3 optimization and the GPU memory usage has been lowered down to 6G compared to 15G without enabling deepseed. However, one strange thing I observed is that the GPU-mem usage will firstly boosts to the normal usage (about 15G), then suddenly reduced to 6G, which means If I increase the batch-size, it will still lead to OOM before the deepspeed optimization started to work. Can you explain why this may happen or how can I found out the reason, since this will not help me to train with larger batchsize.
(3) One possible reason I guess about (2) is the input activations which use much more memory than the model parameters, so I tried the activation checkpointing provided by deepspeed. However, following errors occurred:
self.netG_engine.backward(l_pix)
File "/root/anaconda3/envs/sr3/lib/python3.7/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
ret_val = func(*args, **kwargs)
File "/root/anaconda3/envs/sr3/lib/python3.7/site-packages/deepspeed/runtime/engine.py", line 1816, in backward
self.optimizer.backward(loss, retain_graph=retain_graph)
File "/root/anaconda3/envs/sr3/lib/python3.7/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
ret_val = func(*args, **kwargs)
File "/root/anaconda3/envs/sr3/lib/python3.7/site-packages/deepspeed/runtime/zero/stage3.py", line 1923, in backward
self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
File "/root/anaconda3/envs/sr3/lib/python3.7/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 62, in backward
scaled_loss.backward(retain_graph=retain_graph)
File "/root/anaconda3/envs/sr3/lib/python3.7/site-packages/torch/tensor.py", line 245, in backward
torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
File "/root/anaconda3/envs/sr3/lib/python3.7/site-packages/torch/autograd/__init__.py", line 147, in backward
allow_unreachable=True, accumulate_grad=True) # allow_unreachable flag
RuntimeError: Function torch::autograd::CopyBackwards returned an invalid gradient at index 1 - expected type TensorOptions(dtype=float, device=cuda:0, layout=Strided, requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt)) but got TensorOptions(dtype=float, device=cpu, layout=Strided, requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))
It seems that some activations/states were actually offloaded to cpu during training, but why it cannot be correctly loaded in backward stage? Everything is OK when I changed the activation-checkpointing to torch.utils.checkpointing. So what happend and how could I fix this error.
The ds_config.json is listed as follows if needed:
"train_batch_size" : 16,
"train_micro_batch_size_per_gpu" : 4,
"optimizer": {
"type": "Adam",
"params": {
"lr": 1e-5,
"betas": [
0.9,
0.999
],
"eps": 1e-8,
"weight_decay": 5e-4
}
},
"scheduler": {
"type": "WarmupLR",
"params": {
"warmup_min_lr": 0,
"warmup_max_lr": 1e-5,
"warmup_num_steps": 10000
}
},
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"offload_param": {
"device": "cpu",
"pin_memory": true
},
"overlap_comm": true,
"contiguous_gradients": true,
"sub_group_size": 1e9,
"allgather_bucket_size": "auto",
"reduce_bucket_size": "auto",
"stage3_prefetch_bucket_size": "auto",
"stage3_param_persistence_threshold": "auto"
},
"activation_checkpointing":{
"partition_activations": false,
"cpu_checkpointing": true,
"contiguous_memory_optimization": false,
"number_checkpoints": null,
"synchronize_checkpoint_boundary": false,
"profile": true
},
"steps_per_print" : 50,
"wall_clock_breakdown" : false
}
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.