CUDA OOM when using "ddp" mode in training

[choieq]

🐛 Bug

When i use ddp, i got this CUDA out of memory error.
But it works when using dp mode in same batch size.. I don't understand this situation.

And why is there a memory stack to GPU 0? Although i don't use GPU 0, There is a lot of memory consumption.

Please reproduce using the BoringModel

 trainer = Trainer(fast_dev_run=False, gpus=args.gpu, max_epochs=args.epoch, distributed_backend='ddp', logger=tb_logger)  # distributed_backend='dp')
 trainer.fit(model=model, train_dataloader=train_loader, val_dataloaders=val_loader)

To Reproduce

Use following BoringModel and post here

Environment

Note: Bugs with code are solved faster ! Colab Notebook should be made public !

• IDE: Please, use our python bug_report_model.py template.

• Colab Notebook: Please copy and paste the output from our environment collection script (or fill out the checklist below manually).

You can get the script and run it with:

wget https://raw.githubusercontent.com/PyTorchLightning/pytorch-lightning/master/tests/collect_env_details.py
# For security purposes, please check the contents of collect_env_details.py before running it.
python collect_env_details.py

• PyTorch Version (e.g., 1.0): 1.6.0
• Pytorch-lightning: 1.2.10
• OS (e.g., Linux): Ubuntu 18.04
• How you installed PyTorch (conda, pip, source): conda
• Python version: 3.x
• CUDA/cuDNN version:10.2

Additional context

[tchaton]

Hey @choieq,

Would it be possible for you to reproduce this behaviour with the BoringModel ? Or share some code ?
We can't help you with more details.

Best,
T.C

[choieq]

Hey @choieq,

Would it be possible for you to reproduce this behaviour with the BoringModel ? Or share some code ?
We can't help you with more details.

Best,
T.C

Hi @tchaton
The code is much relative internal codes.. It is hard to translate to colab.
I can share only model definition code.

https://colab.research.google.com/drive/19w29h2MEGbmdX6kcSpaBSPP_IsD-bf9J?usp=sharing
In the last cell, there is BoringModel class.

Any some idea is would help me.

Thanks,

[sualehasif]

I am having an identical extremely annoying issue when using gptneo-1.3b from huggingface on 8 A100s with ddp turned on. This means that I can't turn on any of the advanced plugins as well (like FairScale or DeepSpeed) since they rely on ddp. (My models seem to train quite fine with just dp)

[sualehasif]

I can try to provide a reproducible example but it seems that with every model I try, "ddp_sharded" and "ddp" massively blow up memory as compared to "dp".

[kaushikb11]

@sualehasif That would be great! Will look into it on priority!

[sualehasif]

@kaushikb11 thank you so much. This would be super helpful if resolved soon.

Here is some code where you can choose to use your favorite dataset:

import csv
import json
import sys
import shutil
import os
import subprocess
import time
import msgpack

import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
from torch import optim
import transformers
from transformers import GPT2TokenizerFast, GPTNeoModel
from torch.utils.tensorboard import SummaryWriter
import pytorch_lightning as pl
from pytorch_lightning import Callback
from pytorch_lightning.plugins import DeepSpeedPlugin
from pytorch_lightning.plugins import DDPPlugin
from pytorch_lightning.plugins import DDPShardedPlugin

hyperparams = dict(
    num_gpus = 7,
    prompts_per_gpu = 1,
    # batches_per_macro = 100,
    batches_per_step = 8,
    k = 1,  # should be >=1 since one of the k is always the right link
    peak_lr = 2e-5,
    warmup_period = 0.05
)

class Tokenizer:
    def __init__(self):
        self.tokenizer = GPT2TokenizerFast.from_pretrained('EleutherAI/gpt-neo-1.3B')
        self.tokenizer.padding_side = "left"
        self.tokenizer.pad_token = self.tokenizer.eos_token

    def __call__(self, texts, pad=True):
        return_tensors = "pt"
        if pad == False:
            return_tensors = None
        tokens = self.tokenizer(
                list(texts),
                add_special_tokens=True,
                truncation=True,
                padding=pad,
                return_tensors=return_tensors,
        )['input_ids']
        return tokens

class Dataset(torch.utils.data.Dataset):
    'Characterizes a dataset for PyTorch'
    def __init__(self, data):
        self.tokenizer = Tokenizer()
        self.data = []

        # TODO: pick your favorite dataset and add it to data.

    def __len__(self):
        'Denotes the total number of samples'
        return 128

    def __getitem__(self, index):
        'Generates one sample of data'
        return self.tokenizer([self.data[index]], pad='max_length').squeeze(0), torch.tensor(index)

class EModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = GPTNeoModel.from_pretrained('EleutherAI/gpt-neo-1.3B')

    def forward(self, tokens):
        hidden_states = self.model(tokens, return_dict=True).last_hidden_state
        embeddings = F.normalize(hidden_states, p=2, dim=1)
        return embeddings

from torch.optim.lr_scheduler import _LRScheduler, warnings

class LinearSchedulerWithWarmup(_LRScheduler):
    def __init__(
        self,
        optimizer: torch.optim.Optimizer,
        num_training_steps: int,
        beg_step: int = 0,
        num_warmup_steps: int = 0,
        last_epoch: int = -1,
        verbose: bool = False
    ):
        if not isinstance(num_training_steps, int) or not num_training_steps >= 0:
            raise ValueError("`num_training_steps` must be an integer greater than 0")

        if not isinstance(num_warmup_steps, int) or not num_warmup_steps >= 0:
            raise ValueError("`num_warmup_steps` must be an integer greater than 0")

        self._num_warmup_steps = num_warmup_steps
        self._num_training_steps = num_training_steps
        self._beg_step = beg_step

        super().__init__(optimizer, last_epoch=last_epoch, verbose=verbose)

    def lr_lambda(self, current_step: int) -> int:

        #assert current_step >= self.beg_step
        assert current_step <= self._num_training_steps

        relative_curr_step = current_step - self._beg_step
        relative_num_training_steps = self._num_training_steps - self._beg_step

        if relative_curr_step < self._num_warmup_steps:
            return float(relative_curr_step) / float(max(1, self._num_warmup_steps))
        return max(
            0.0,
            float(relative_num_training_steps - relative_curr_step) /
            float(max(1, relative_num_training_steps - self._num_warmup_steps))
        )

    def get_lr(self):
        if not self._get_lr_called_within_step:
            warnings.warn("To get the last learning rate computed by the scheduler, " "please use `get_last_lr()`.")

        return [base_lr * self.lr_lambda(self.last_epoch) for base_lr in self.base_lrs]


class TrainableModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.completion_model = EModel()

    # not used during training, only during inference
    def forward(self, batch):
        return self.completion_model(batch)

    def training_step(self, batch, batch_idx):
        prompts_tokens, link_idxs = batch
        query_vectors = self.completion_model(prompts_tokens)

        return query_vectors

    def training_step_end(self, batch_parts):
        v = batch_parts
        logits = torch.matmul(v, v.T)

        labels = torch.arange(len(v)).to(self.device)
        loss = F.cross_entropy(logits, labels)

        return loss

    def configure_optimizers(self):
        optimizer = transformers.AdamW(self.parameters(), lr=hyperparams['peak_lr'])
        scheduler = LinearSchedulerWithWarmup(
            optimizer,
            num_warmup_steps=int(hyperparams['warmup_period'] * hyperparams["num_gradient_updates"]),
            num_training_steps=hyperparams["num_gradient_updates"],
        )

        return {
            'optimizer': optimizer,
            'lr_scheduler': {
                'scheduler': scheduler,
                'interval': 'step',
                'frequency': 1,
            }
        }

class CUDACallback(Callback):
    def on_train_epoch_start(self, trainer, pl_module):
        # Reset the memory use counter
        torch.cuda.reset_peak_memory_stats(trainer.root_gpu)
        torch.cuda.synchronize(trainer.root_gpu)
        self.start_time = time.time()

    def on_train_epoch_end(self, trainer, pl_module, outputs):
        torch.cuda.synchronize(trainer.root_gpu)
        max_memory = torch.cuda.max_memory_allocated(trainer.root_gpu) / 2 ** 20
        epoch_time = time.time() - self.start_time

        max_memory = torch.tensor(max_memory, dtype=torch.int, device=trainer.root_gpu)
        epoch_time = torch.tensor(epoch_time, dtype=torch.int, device=trainer.root_gpu)

        torch.distributed.all_reduce(max_memory, op=torch.distributed.ReduceOp.SUM)
        torch.distributed.all_reduce(epoch_time, op=torch.distributed.ReduceOp.SUM)

        world_size = torch.distributed.get_world_size()

        print(f"Average Epoch time: {epoch_time.item() / float(world_size):.2f} seconds")
        print(f"Average Peak memory {max_memory.item() / float(world_size):.2f}MiB")

# TODO: initialize the dataset with data
dataset=Dataset()

dataloader = torch.utils.data.DataLoader(dataset, 
                                        batch_size=hyperparams['num_gpus'], 
                                        shuffle=True,
                                        num_workers=os.cpu_count())

sharded_plugin = DDPShardedPlugin()
model = TrainableModel()
trainer = pl.Trainer(
    gpus=hyperparams['num_gpus'],
    num_nodes=1,
    accelerator='ddp',
    accumulate_grad_batches=hyperparams['batches_per_step'],
    log_every_n_steps=1,
    callbacks=[CUDACallback()],
    progress_bar_refresh_rate=1,
    max_epochs=1,
    profiler="simple",
    precision=16,
    plugins=DDPShardedPlugin,
    # plugins='deepspeed_stage_2'
)

[kaushikb11]

@sualehasif @choieq Something to be noted that the effective batch size for ddp and ddp based plugins is batch_size x num_gpus. That could be the issue for OOM.

[sualehasif]

Yep. I know that but there seems to be no way to avoid that? For e.g. in my example, I balanced the effective batch_size in ddp to the batch_size in dp and it didnt seem to help.

@kaushikb11 If you have a way to run a large model like GPT-NeoX using deepspeed or ddp_sharded, feel free to post a script.

[kaushikb11]

Sure, you could run the minGPT model on several plugins and benchmark

[sualehasif]

I can't run minGPT with GPT-Neo-1.3B. It gives the same OOM errors even on 8A100s.