deepspeedai / deepspeedai/DeepSpeed

OOM while llama2-70B SFT

Open
#5,228 1 comment 3 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug training
Dominant language
Python
Stars
43.1k
Forks
5k
Avg merge
4d 15h
Merged PRs (30d)
112

Description

Describe the bug

System Info
H100X8 (80GB)
docker : nvcr.io/nvidia/pytorch:23.12-py3
python packages:
  datasets 
  evaluate
  accelerate=
  transformers
  deepspeed

getting OOM when doing Supervised Fine tuning for llama2-70B using zero-3 even with batch size of 1.

To Reproduce
Steps to reproduce the behavior:
Script

import argparse
import os
import torch
import transformers
from accelerate import Accelerator
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaForCausalLM
import numpy as np
import datasettokenizer as tok
import time
import math
import deepspeed
from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
from deepspeed.accelerator import get_accelerator

accelerator = Accelerator()

parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', type=int, required=False, default=0)

args = parser.parse_args()

model_name = "meta-llama/Llama-2-70b-hf"
output_dir = "tmp"
max_length = 512
max_steps = 2000
learning_rate = 2.5e-5
learning_rate = 0.00002
batch_size=8
num_workers=8
epochs=10

train_ds_size = 1000
steps_per_epoch = train_ds_size // (batch_size * num_workers)

world_size = math.ceil(num_workers // 8)

print("Loading tiny_shakespeare dataset")
dataset = load_dataset("tiny_shakespeare") #,streaming=True)

tokenizer = AutoTokenizer.from_pretrained(
    model_name,
    padding_side="left",
    add_eos_token=True,
    add_bos_token=True,
)
tokenizer.pad_token = tokenizer.eos_token

dataset = tok.tokenize_dataset(dataset, tokenizer, block_size=max_length)

tokenized_train_dataset = dataset["train"]
tokenized_val_dataset = dataset["validation"]

deepspeed_cfg = {
    "fp16": {
        "enabled": True,
        "initial_scale_power": 8,
    },
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": learning_rate,
            "betas": [0.9, 0.999],
            "eps": 1e-8,
        },
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": learning_rate,
            "warmup_max_lr": learning_rate,
            "warmup_num_steps": 0
        },
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": True,
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": True,
        },
        "overlap_comm": True,
        "contiguous_gradients": True,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "gather_16bit_weights_on_model_save": True,
        "round_robin_gradients": True,
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": 1.0,
    "steps_per_print": 10,
    "train_batch_size":batch_size * 1 *num_workers,
    "train_micro_batch_size_per_gpu":batch_size,
    "wall_clock_breakdown": False,
}

print("Loading model")

model = MixtralForCausalLM.from_pretrained(model_name, use_cache=False, low_cpu_mem_usage=True, torch_dtype=torch.float16)
model.resize_token_embeddings(len(tokenizer))

model, _, _, _ = deepspeed.initialize(model=model, config=deepspeed_cfg)

print("Model loaded")

model.train()
for epoch in range(epochs):
    for i, d in enumerate(train_loader):
        if i > steps_per_epoch:
            break
        iid = d['input_ids'].cuda()
        am = d['attention_mask'].cuda()
        outputs = model(input_ids=iid,attention_mask=am)
        loss = compute_loss(outputs, d['labels'].cuda())
        model.backward(loss)
        model.step()
        if torch.distributed.get_rank() == 0:
            print(f"iter: {i} loss:, {loss.item()}")

Error

Loading checkpoint shards: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 15/15 [00:01<00:00, [173/1849]
Loading checkpoint shards: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 15/15 [00:01<00:00,  8.88it/s]
[2024-03-05 22:00:14,149] [INFO] [logging.py:96:log_dist] [Rank -1] DeepSpeed info: version=0.13.1, git-hash=unknown, git-branch=unknown                                                                              
[2024-03-05 22:00:14,149] [INFO] [comm.py:637:init_distributed] cdb=None                                                                                                                                              
[2024-03-05 22:00:14,173] [INFO] [logging.py:96:log_dist] [Rank -1] DeepSpeed info: version=0.13.1, git-hash=unknown, git-branch=unknown                                                                              
[2024-03-05 22:00:14,173] [INFO] [comm.py:637:init_distributed] cdb=None                                                                                                                                              
[2024-03-05 22:00:14,177] [INFO] [logging.py:96:log_dist] [Rank -1] DeepSpeed info: version=0.13.1, git-hash=unknown, git-branch=unknown                                                                              
[2024-03-05 22:00:14,177] [INFO] [comm.py:637:init_distributed] cdb=None                                                                                                                                              
Traceback (most recent call last):                                                                                                                                                                                    
  File "/workspace/llama-ft/t1.py", line 106, in <module>                                                                                                                                                             
    model, _, _, _ = deepspeed.initialize(model=model, config=deepspeed_cfg)                                                                                                                                          
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/__init__.py", line 171, in initialize                                                                                                                       
Traceback (most recent call last):                                                                                                                                                                                    
  File "/workspace/llama-ft/t1.py", line 106, in <module>                                                                                                                                                             
    engine = DeepSpeedEngine(args=args,                                                                                                                                                                               
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/engine.py", line 263, in __init__                                                                                                                   
    self._configure_distributed_model(model)                                                                                                                                                                          
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/engine.py", line 1103, in _configure_distributed_model                                                                                              
    model, _, _, _ = deepspeed.initialize(model=model, config=deepspeed_cfg)                                                                                                                                          
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/__init__.py", line 171, in initialize
    engine = DeepSpeedEngine(args=args,
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/engine.py", line 263, in __init__
    self.module.to(self.device)
  File "/usr/local/lib/python3.10/dist-packages/transformers/modeling_utils.py", line 2597, in to
    self._configure_distributed_model(model)
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/engine.py", line 1103, in _configure_distributed_model
    self.module.to(self.device)
  File "/usr/local/lib/python3.10/dist-packages/transformers/modeling_utils.py", line 2597, in to
    return super().to(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1151, in to
    return super().to(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1151, in to
    return self._apply(convert)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 801, in _apply
    return self._apply(convert)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 801, in _apply
    module._apply(fn)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 801, in _apply
    module._apply(fn)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 801, in _apply
    module._apply(fn)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 801, in _apply
    module._apply(fn)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 801, in _apply
    module._apply(fn)
      [Previous line repeated 2 more times]
module._apply(fn)  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 824, in _apply

  [Previous line repeated 2 more times]
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 824, in _apply
        param_applied = fn(param)param_applied = fn(param)

  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1149, in convert
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1149, in convert
    return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 448.00 MiB. GPU 7 has a total capacity of 79.14 GiB of which 398.81 MiB is free. Process 3147356 has 78.75 GiB memory in use. Of the allocated memo
ry 78.24 GiB is allocated by PyTorch, and 487.50 KiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation.  See documentation for 
Memory Management and PYTORCH_CUDA_ALLOC_CONF

Expected behavior
should work without OOM

ds_report output

[2024-03-05 22:04:18,143] [INFO] [real_accelerator.py:191:get_accelerator] Setting ds_accelerator to cuda (auto detect)
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
fused_adam ............. [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_lion ............... [NO] ....... [OKAY]
 [WARNING]  Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH
evoformer_attn ......... [NO] ....... [NO]
fused_lamb ............. [NO] ....... [OKAY]
fused_lion ............. [NO] ....... [OKAY]
inference_core_ops ..... [NO] ....... [OKAY]
cutlass_ops ............ [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
ragged_device_ops ...... [NO] ....... [OKAY]
ragged_ops ............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.2
 [WARNING]  using untested triton version (2.1.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/usr/local/lib/python3.10/dist-packages/torch']
torch version .................... 2.2.0a0+81ea7a4
deepspeed install path ........... ['/usr/local/lib/python3.10/dist-packages/deepspeed']
deepspeed info ................... 0.13.1, unknown, unknown
torch cuda version ............... 12.3
torch hip version ................ None
nvcc version ..................... 12.3
deepspeed wheel compiled w. ...... torch 2.2, cuda 12.3
shared memory (/dev/shm) size .... 1007.72 GB

Screenshots
If applicable, add screenshots to help explain your problem.

System info (please complete the following information):

  • OS: [e.g. Ubuntu 18.04] : Ubuntu 20.4
  • GPU count and types [e.g. two machines with x8 A100s each] 1 machine with x8 H100s each
  • Interconnects (if applicable) [e.g., two machines connected with 100 Gbps IB]
  • Python version- 3.10
  • Any other relevant info about your setup

Launcher context
Are you launching your experiment with the deepspeed launcher, MPI, or something else?

Docker context
Are you using a specific docker image that you can share?

Additional context
Add any other context about the problem here.

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.

Research direction

Start with the model initialization at line 106 of t1.py and the shown DeepSpeed configuration, then review the ds_report output and the CUDA out-of-memory traceback. Reproduce the failure with the provided Llama-2-70B SFT setup and verify that the same eight-GPU configuration initializes and trains without OOM.

Written by the indexing model from the issue text.

Assessment

Tech stack
docker, python, pytorch
Domain
distributed-systems, machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.