Deepspeed: gpt-neo-2.7B not trackable with RTX 3090 and 64GB RAM?

[neil-tan]

Hello there,
I'm having trouble fine-tuning HF pre-trained transformer EleutherAI/gpt-neo-2.7B on my RTX 3090 with 64GB RAM. With stage=3 off-loading, the RAM usage seems excessively high. Is it normal to run out of 64GB of system memory with Adam optimizer for a 2.7B parameters model?

Error Short

; cpu off-loading
[2022-07-10 10:53:21,611] [INFO] [utils.py:829:see_memory_usage] MA 10.75 GB         Max_MA 10.75 GB         CA 10.76 GB         Max_CA 17 GB 
[2022-07-10 10:53:21,612] [INFO] [utils.py:837:see_memory_usage] CPU Virtual Memory:  used = 59.96 GB, percent = 95.6%
[2022-07-10 10:53:22,068] [INFO] [utils.py:828:see_memory_usage] before backward
[2022-07-10 10:53:22,069] [INFO] [utils.py:829:see_memory_usage] MA 11.93 GB         Max_MA 12.39 GB         CA 12.43 GB         Max_CA 12 GB 
[2022-07-10 10:53:22,069] [INFO] [utils.py:837:see_memory_usage] CPU Virtual Memory:  used = 59.98 GB, percent = 95.7%
[2022-07-10 10:53:22,177] [INFO] [utils.py:828:see_memory_usage] before optimizer
[2022-07-10 10:53:22,178] [INFO] [utils.py:829:see_memory_usage] MA 11.91 GB         Max_MA 11.93 GB         CA 12.43 GB         Max_CA 12 GB 
[2022-07-10 10:53:22,178] [INFO] [utils.py:837:see_memory_usage] CPU Virtual Memory:  used = 59.98 GB, percent = 95.7%
Killed

Complete Logs
cpu off-loading
nvme off-loading

System & Environment
RTX 3090, CUDA 11.6
RAM 64GB installed
Torch 1.12.0
Pytorch-lightning 1.6.4
Deep-speed 0.6.5
Python 3.8.6
pip freeze

Code

from argparse import ArgumentParser
import torch
import transformers
from transformers import AutoTokenizer
import pytorch_lightning as pl
from pytorch_lightning import seed_everything
from lightning_transformers.utilities.deepspeed import enable_transformers_pretrained_deepspeed_sharding
from lightning_transformers.task.nlp.language_modeling import (
    LanguageModelingDataModule,
    LanguageModelingTransformer,
)
import deepspeed
from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live
from deepspeed.runtime.utils import see_memory_usage
from pytorch_lightning.strategies import DeepSpeedStrategy

seed_everything(8653745, workers=True)

train_config = {
  'lr': 0.001,
  'wd': 5e-4,
  "beta": (0.9, 0.999),
  "batch_size": 1,
}

parser = ArgumentParser(conflict_handler='resolve')
parser = pl.Trainer.add_argparse_args(parser)
parser.add_argument("--gpus", default=1)
parser.add_argument("--max_epochs", default=1)
parser.add_argument("--deterministic", default=True)
parser.add_argument("--strategy", default=DeepSpeedStrategy(
                                                          stage=3,
                                                          offload_optimizer=True,
                                                          offload_parameters=True,
                                                          params_buffer_size = 150_000_000,
                                                          logging_level="INFO",
                                                          remote_device="nvme",
                                                          offload_optimizer_device="nvme",
                                                          offload_params_device="nvme",
                                                          nvme_path="/home/neil/tmp/deepspeed_offloading",
                                                      ))
# parser.add_argument("--precision", default=16)
config_all = parser.parse_args("")
train_config.update(vars(config_all))

model = transformers.GPTNeoForCausalLM.from_pretrained('EleutherAI/gpt-neo-2.7B')

tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="EleutherAI/gpt-neo-2.7B")

dm = LanguageModelingDataModule(
    batch_size=1,
    dataset_name="wikitext",
    dataset_config_name="wikitext-2-raw-v1",
    max_length=model.config.max_position_embeddings,
    tokenizer=tokenizer,
    num_workers=20,
)

class DummyLMModel(pl.LightningModule):
  def __init__(self, **config):
    super().__init__()
    self.save_hyperparameters(config)

  def setup(self, stage=None) -> None:
    # https://lightning-transformers.readthedocs.io/en/latest/features/large_model_training.html
    if not hasattr(self, "model"):
        enable_transformers_pretrained_deepspeed_sharding(self)
        self.model = model
        estimate_zero3_model_states_mem_needs_all_live(self.model, num_gpus_per_node=1, num_nodes=1, additional_buffer_factor=1.5)
        see_memory_usage(f'after setup', force=True)

  def forward(self, **inputs):
    return self.model(**inputs)

  def configure_optimizers(self):
    return deepspeed.ops.adam.DeepSpeedCPUAdam(
    # return deepspeed.ops.adam.FusedAdam(
                            self.model.parameters(),
                            lr=self.hparams.lr, betas=self.hparams.beta,
                            weight_decay=self.hparams.wd)

  def training_step(self, batch, batch_idx):
    see_memory_usage(f'before forward', force=True)
    output = self(input_ids=batch['input_ids'], attention_mask=batch['attention_mask'], labels=batch['labels'])
    see_memory_usage(f'before backward', force=True)
    return output.loss

  def training_step_end(self, training_step_outputs):
    see_memory_usage(f'before optimizer', force=True)
    return {'loss': training_step_outputs}


pl_model = DummyLMModel(**train_config)

trainer = pl.Trainer.from_argparse_args(config_all)
trainer.fit(pl_model, dm)

[Syntax3rror404]

To handle this big models i recommend using deepspeed for fix high ram usage.