deepspeedai / deepspeedai/DeepSpeed

Still OOM when enabling deepspeed and error during backward

Open
#3,313 2 comments 0 reactions 1 assignee View on GitHub

@mrwyattii is already working on this.

Since Apr 21, 2023.

training
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.