GPU experiments - e.g. GTX 2080Ti

[tomekrut]

I have been doing some tests on GTX 2080Ti with a simple LSTM network

• 1x GPU - 84M parameters is the max I was able to put in without DeepSpeed
• 2x GPUs including DeepSpeed... 105M parameters... 20M more parameters
• 4x GPUs including DeepSpeed... 135M parameters
  so, roughly 60M parameters are fixed on each device and the positive delta is 20M...

Can you share your experience as well? Below is the code that you can copy paste and run in as it is. I simply changed the HIDDEN_DIM to play with the number of parameters and then either

python experiment.py
or

deepspeed experiment.py --deepspeed --deepspeed_config ds_config.json

File - ds_config.json

{
  "train_batch_size": 128,
  "steps_per_print": 2000,
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 0.001,
      "betas": [
        0.8,
        0.999
      ],
      "eps": 1e-8,
      "weight_decay": 3e-7
    }
  },
  "scheduler": {
    "type": "WarmupLR",
    "params": {
      "warmup_min_lr": 0,
      "warmup_max_lr": 0.001,
      "warmup_num_steps": 1000
    }
  },
  "wall_clock_breakdown": false
}

File - experiment.py

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data.dataset import Dataset
import deepspeed as ds
import argparse
import datetime
import torch.optim as optim

def add_argument():
    parser = argparse.ArgumentParser(description='MyCustom')
    parser.add_argument('--with_cuda', default=False, action='store_true',
                        help='use CPU in case there\'s no GPU support')
    parser.add_argument('--use_ema', default=False, action='store_true',
                        help='whether use exponential moving average')
    parser.add_argument('-b', '--batch_size', default=32, type=int,
                        help='mini-batch size (default: 32)')
    parser.add_argument('-e', '--epochs', default=30, type=int,
                        help='number of total epochs (default: 30)')
    parser.add_argument('--local_rank', type=int, default=-1,
                        help='local rank passed from distributed launcher')
    parser = ds.add_config_arguments(parser)
    args = parser.parse_args()
    return args

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


class RNN(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim, n_layers,
                 bidirectional, dropout, pad_idx):
        super().__init__()

        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=pad_idx)
        self.rnn = nn.LSTM(embedding_dim,
                           hidden_dim,
                           num_layers=n_layers,
                           bidirectional=bidirectional,
                           dropout=dropout)
        self.fc = nn.Linear(hidden_dim * 2, output_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, text, text_lengths):
        embedded = self.dropout(self.embedding(text))
        packed_embedded = nn.utils.rnn.pack_padded_sequence(embedded, text_lengths)
        packed_output, (hidden, cell) = self.rnn(packed_embedded)
        hidden = self.dropout(torch.cat((hidden[-2, :, :], hidden[-1, :, :]), dim=1))
        return self.fc(hidden)


SEED = 1234
torch.manual_seed(SEED)
torch.backends.cudnn.deterministic = True


def generate_data():
    d, l = [], []
    for r in range(1000):
        v = np.random.choice(25000, 512)
        d.append(v)
        l.append(float(np.random.randint(0, 2)))
    return np.array(d), np.array(l)


class MyCustomDataset(Dataset):
    def __init__(self):
        self.data, self.label = generate_data()

    def __getitem__(self, index):
        return (self.data[index], self.label[index])

    def __len__(self):
        return self.data.shape[0]


BATCH_SIZE = 64
INPUT_DIM = 25000
EMBEDDING_DIM = 100
HIDDEN_DIM = 1350
OUTPUT_DIM = 1
N_LAYERS = 3
BIDIRECTIONAL = True
DROPOUT = 0.5
PAD_IDX = 1

train_data = MyCustomDataset()
train_loader = torch.utils.data.DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)
model = RNN(INPUT_DIM,
            EMBEDDING_DIM,
            HIDDEN_DIM,
            OUTPUT_DIM,
            N_LAYERS,
            BIDIRECTIONAL,
            DROPOUT,
            PAD_IDX)

print(f'The model has \x1b[6;30;42m {count_parameters(model):,} \x1b[0m trainable parameters')
criterion = nn.BCEWithLogitsLoss()
args = add_argument()

if args.deepspeed:
    parameters = filter(lambda p: p.requires_grad, model.parameters())
    model_engine, optimizer, train_loader, _ = ds.initialize(args=args, model=model, model_parameters=parameters, training_data=train_data)

    for epoch in range(1):
        running_loss = 0.0
        start = datetime.datetime.now()
        for i, data in enumerate(train_loader):
            inputs, labels = data[0].to(model_engine.local_rank).T, data[1].to(model_engine.local_rank)

            text_lengths = torch.ones(inputs.shape[1])
            text_lengths = text_lengths.fill_(512)
            text_lengths = text_lengths.to(model_engine.local_rank)

            outputs = model_engine(inputs, text_lengths).squeeze(1)
            loss = criterion(outputs, labels)

            model_engine.backward(loss)
            model_engine.step()

            running_loss += loss.item()

        end = datetime.datetime.now() - start
        print(f'Single epoch time {end}')

else:
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model_engine = model.to(device)
    optimizer = optim.Adam(model_engine.parameters())
    for epoch in range(1):
        running_loss = 0.0
        start = datetime.datetime.now()
        for i, data in enumerate(train_loader):
            optimizer.zero_grad()

            inputs, labels = data[0].to(device).T, data[1].to(device)

            text_lengths = torch.ones(inputs.shape[1])
            text_lengths = text_lengths.fill_(512)
            text_lengths = text_lengths.to(device)

            outputs = model_engine(inputs, text_lengths).squeeze(1)
            loss = criterion(outputs, labels)

            loss.backward()

            optimizer.step()

        end = datetime.datetime.now() - start
        print(f'Single epoch time {end}')

[tjruwase]

Thanks for sharing this experiment. I do have one question. Are you keeping the effective batch size constant regardless of the GPU count? I notice train_batch_size is 128 in the json config, but that BATCH_SIZE is hard coded to 64 in the script.

[tomekrut]

Looking at the paper again:

Without any model-parallelism,ZeRO-OS-powereddata-parallelism enables training mod-els with up to 6B parameters on V100 GPUs (32GB of device memory) while existingsystems (e.g., PyTorch Distributed Data Parallel) runs out of memory with 1.5B param-eter models.ZeRO-OSgives 4x memory saving / model size boost. Model developerscan run up to 6B parameter models without worrying about model parallelism

Assuming that the authors talk about DGX-2 which has 16 devices then:

• additional 4.5B / 15 devices = 0.3B gain per device
• 1.5B on a single device / 0.3B = effective gain is 20% per device

In my above case is 80M / 20M ~ 25% gain.. I guess it would make sense then?

(on the other hand, if I understand correctly) Table 1 then shows that for 7.5B model only 4 devices are needed (DP=4) as they consume 30GB of memory? It shows nice linear growth.

If this one is then correct... why I am getting such a small improvement? I should be getting 80M x 4GPUs = 320M parameters.

[tomekrut]

Thanks for sharing this experiment. I do have one question. Are you keeping the effective batch size constant regardless of the GPU count? I notice train_batch_size is 128 in the json config, but that BATCH_SIZE is hard coded to 64 in the script.

I was trying to do the following:

• no DeepSpeed then 64
• with DeepSpeed 128 then train_micro_batch_size_per_gpu will be 64 when I use 2x GPUs... this was my default setting.

[tjruwase]

I was trying to do the following:

• no DeepSpeed then 64
• with DeepSpeed 128 then train_micro_batch_size_per_gpu will be 64 when I use 2x GPUs... this was my default setting.

Thanks for the clarification. I assume for 4x GPUs setting, you set train_batch_size to 256 in the json config.

[tomekrut]

I was trying to do the following:

• no DeepSpeed then 64
• with DeepSpeed 128 then train_micro_batch_size_per_gpu will be 64 when I use 2x GPUs... this was my default setting.

Thanks for the clarification. I assume for 4x GPUs setting, you set train_batch_size to 256 in the json config.

Well, I missed this one... but in advantage to your framework :) If I had increased to 256 I wouldn't have been able to probably even achieve 135M but less... don't have 4 available now to check as they are being utilized now.

[tjruwase]

Well, I missed this one... but in advantage to your framework :) If I had increased to 256 I wouldn't have been able to probably even achieve 135M but less... don't have 4 available now to check as they are being utilized now.

Hahaha :). Got it. Please do share your experience when you get the chance. I will address your other questions on the expected gains separately.

[ShadenSmith]

The code above only uses data parallelism, and so the full model, gradients, and optimizer states are replicated on each process. You can save memory by enabling ZeRO. In its current implementation (Stage 1), the optimizer states are partitioned across processes. Further ZeRO optimization implementations are a work in progress. We have a summary of some upcoming improvements here: https://www.deepspeed.ai/news/

To save additional memory, with some models you can also use model parallelism which splits the model itself (and gradients, etc.) across devices. DeepSpeed works with model parallelism if the user provides it. We have more details on model parallelism + DeepSpeed in the Megatron tutorial.

Note that the table showing 30GB assumes all three stages of ZeRO are implemented. With only Stage 1, you need 64 V100s for the 7.5B parameter model.

[tomekrut]

@ShadenSmith , great thanks... I will try updating the numbers soon.

[tomekrut]

@ShadenSmith @jeffra - It is a bit unrelated to the main topic of this issue but doing some performance analysis - I tried few LSTMs with zero_optimization true and each time I get

[deepspeed] OVERFLOW! Skipping step. Attempted loss scale:

Have I missed something?

e.g. the simplest possible

class LSTMer(nn.Module):

    def __init__(self, vocab_size, embedding_dim, hidden_dim, tagset_size):
        super(LSTMer, self).__init__()
        self.hidden_dim = hidden_dim
        self.word_embeddings = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim)
        self.hidden2tag = nn.Linear(hidden_dim, tagset_size)

    def forward(self, sentence):
        embeds = self.word_embeddings(sentence)
        self.lstm.flatten_parameters()
        lstm_out, _ = self.lstm(embeds)
        tag_space = self.hidden2tag(lstm_out)
        tag_scores = F.log_softmax(tag_space, dim=1)
        return tag_scores[:, -1].squeeze(1)

[ShadenSmith]

Can you share more of your logs? Does it continue to always overflow, or does the loss scaling eventually decrease and training proceeds?

Some number of overflow steps is typical. DeepSpeed decreases the scaling each time it occurs and it should eventually get out of an overflow state and begin doing useful training.

[tomekrut]

Here is the exact plug and play script that you can run including the config file.
I have seen two patterns:

• it finished but not sure about the quality of the model as the script was created for the purpose of checking how many parameters I could fit into the GPUs,
• started training but ran out of memory eventually - when the number of parameters was defined close to the maximum of given GPUs capacity,

{
  "train_batch_size": 16,
  "zero_optimization": true,
  "steps_per_print": 2000,
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 0.001,
      "betas": [
        0.8,
        0.999
      ],
      "eps": 1e-8,
      "weight_decay": 3e-7
    }
  },
  "scheduler": {
    "type": "WarmupLR",
    "params": {
      "warmup_min_lr": 0,
      "warmup_max_lr": 0.001,
      "warmup_num_steps": 1000
    }
  },
  "wall_clock_breakdown": false,
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "initial_scale_power": 32,
    "loss_scale_window": 1000,
    "hysteresis": 2,
 	"min_loss_scale": 1
    }
}

sample script

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data.dataset import Dataset
import torch.nn.functional as F
import deepspeed as ds
import argparse
import datetime
import torch.optim as optim


def add_argument():
    parser = argparse.ArgumentParser(description='MyCustom')
    parser.add_argument('--with_cuda', default=False, action='store_true',
                        help='use CPU in case there\'s no GPU support')
    parser.add_argument('--use_ema', default=False, action='store_true',
                        help='whether use exponential moving average')
    parser.add_argument('-b', '--batch_size', default=32, type=int,
                        help='mini-batch size (default: 32)')
    parser.add_argument('-e', '--epochs', default=30, type=int,
                        help='number of total epochs (default: 30)')
    parser.add_argument('--local_rank', type=int, default=-1,
                        help='local rank passed from distributed launcher')
    parser = ds.add_config_arguments(parser)
    args = parser.parse_args()
    return args


def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


class LSTMer(nn.Module):

    def __init__(self, vocab_size, embedding_dim, hidden_dim, tagset_size):
        super(LSTMer, self).__init__()
        self.hidden_dim = hidden_dim
        self.word_embeddings = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim)
        self.hidden2tag = nn.Linear(hidden_dim, tagset_size)

    def forward(self, sentence):
        embeds = self.word_embeddings(sentence)
        self.lstm.flatten_parameters()
        lstm_out, _ = self.lstm(embeds)
        tag_space = self.hidden2tag(lstm_out)
        tag_scores = F.log_softmax(tag_space, dim=1)
        return tag_scores[:, -1].squeeze(1)


SEED = 1234
torch.manual_seed(SEED)
torch.backends.cudnn.deterministic = True


def generate_data():
    d, l = [], []
    for r in range(500):
        v = np.random.choice(25000, 512)
        d.append(v)
        l.append(float(np.random.randint(0, 2)))
    return np.array(d), np.array(l)


class MyCustomDataset(Dataset):
    def __init__(self):
        self.data, self.label = generate_data()

    def __getitem__(self, index):
        return (self.data[index], self.label[index])

    def __len__(self):
        return self.data.shape[0]


BATCH_SIZE = 8
INPUT_DIM = 25000
EMBEDDING_DIM = 350
HIDDEN_DIM = 5000
OUTPUT_DIM = 1
N_LAYERS = 3

train_data = MyCustomDataset()
train_loader = torch.utils.data.DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)

model = LSTMer(INPUT_DIM,
               EMBEDDING_DIM,
               HIDDEN_DIM,
               OUTPUT_DIM)

print(f'The model has \x1b[6;30;42m {count_parameters(model):,} \x1b[0m trainable parameters')
criterion = nn.BCEWithLogitsLoss()
args = add_argument()

if args.deepspeed:
    parameters = filter(lambda p: p.requires_grad, model.parameters())
    model_engine, optimizer, train_loader, _ = ds.initialize(args=args, model=model, model_parameters=parameters, training_data=train_data)

    for epoch in range(1):
        running_loss = 0.0
        start = datetime.datetime.now()
        for i, data in enumerate(train_loader):
            inputs, labels = data[0].to(model_engine.local_rank), data[1].to(model_engine.local_rank)

            outputs = model_engine(inputs)
            loss = criterion(outputs, labels.half())

            model_engine.backward(loss)
            model_engine.step()

            running_loss += loss.item()

else:
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model_engine = model.to(device)
    optimizer = optim.Adam(model_engine.parameters())
    for epoch in range(1):
        running_loss = 0.0
        start = datetime.datetime.now()
        for i, data in enumerate(train_loader):
            optimizer.zero_grad()

            inputs, labels = data[0].to(device), data[1].to(device)
            outputs = model_engine(inputs)
            loss = criterion(outputs, labels)

            loss.backward()

            optimizer.step()

end = datetime.datetime.now() - start
print(f'Single epoch time {end}')

[loadams]

Closing this issue as it is stale and the current cuda/torch versions do not support the 2080Ti anymore. If you are having any issues, please open a new issue and link this and we can take a look. Thanks for using DeepSpeed!