# Accelerate/DeepSpeed: Flan-T5 OOM despite device\_mapping

**URL:** </static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Faccelerate-deepspeed-flan-t5-oom-despite-device-mapping%2F27764%3E%5C
**Category:** 🤗Accelerate\
**Created:** [December 13, 2022, 4:41pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Faccelerate-deepspeed-flan-t5-oom-despite-device-mapping%2F27764 "2022-12-13T16:41:46Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![Breenori](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/breenori/32/13027_2.png) [@Breenori](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FBreenori)%5C
**Post date:** [December 13, 2022, 4:41pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Faccelerate-deepspeed-flan-t5-oom-despite-device-mapping%2F27764%2F1 "2022-12-13T16:41:46Z")

</div>

I currently want to get FLAN-T5 working for inference on my setup which consists of 6x RTX 3090 (6x. 24GB) and cannot get it to work in my Jupyter Notebook inside a Pytorch Nvidia Container (22.06).

I have already tried configuring DeepSpeed and Accelerate in order to reduce the size of the model and to distribute it over all GPUs. This works fine at the start, but only allocates about 10GB on every GPU. This is followed by GPU 0 being loaded fully and eventually encountering an OOM error.

 ![FLAN-T5-OOM](https://us1.discourse-cdn.com/hellohellohello/original/2X/8/8c166cac1b1617cf98eec4a2ac0adc343fe3c2dd.png)

This is the code I use to initialize accelerate and DeepSpeed.

```auto
# Set environment variable to store the file locally to avoid redownloading.
import os
!mkdir -p /home/jovyan/datasets/huggingface
os.environ['TRANSFORMERS_CACHE'] = '/home/jovyan/datasets/huggingface'

from transformers import T5Tokenizer, T5ForConditionalGeneration
import torch
import accelerate

# Prepare deepspeed
deepspeed = accelerate.DeepSpeedPlugin(gradient_accumulation_steps=1,
                                       zero3_init_flag=False,
                                       zero_stage=3,
                                       offload_optimizer_device="cpu",
                                       offload_param_device="cpu")
accelerator = accelerate.Accelerator(deepspeed_plugin=deepspeed)
accelerate.state.AcceleratorState().deepspeed_plugin.deepspeed_config['train_micro_batch_size_per_gpu'] = 1
accelerate.state.AcceleratorState()

```

Afterwards I try to load the models using accelerator.prepare, but here I run into the problem:

```auto
model, tokenizer = accelerator.prepare(
    T5ForConditionalGeneration.from_pretrained("google/flan-t5-xxl", device_map="balanced_low_0"), 
    T5Tokenizer.from_pretrained("google/flan-t5-xxl", device_map="balanced_low_0")
)

```

I already tried all Zero stages, with and out deepspeed, with and without param/optimizer offloading and am out of ideas. I also tried setting num\_processes but ran into an issue where I don’t know where to set acc\_gradient\_steps or world\_size and can’t find any documentation either.

I’m sorry if this is a mundane task, but any help would be greatly appreciated.

---

<div class="post-metadata">

**Author:** ![smangrul](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/smangrul/32/3791_2.png) [@smangrul](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fsmangrul)%5C
**Post date:** [January 2, 2023, 9:56am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Faccelerate-deepspeed-flan-t5-oom-despite-device-mapping%2F27764%2F2 "2023-01-02T09:56:18Z")

</div>

Hello @Breenori, `device_mapping` is supported only for inference. Could you try removing it and then use DeepSpeed Stage 3 with CPU offloading?