

TITLE
[deepspeed] `bigscience/T0*` multi-gpu inference with ZeRO

COMMENTS
32

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
## Environment info

- `transformers` version: 4.17.0.dev0
- Platform: Linux-5.13.0-27-generic-x86_64-with-glibc2.10
- Python version: 3.8.0
- PyTorch version (GPU?): 1.10.1 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: yes (deepspeed)
- Note: I installed DeepSpeed from source

### Who can help

Models:
(I'm actually trying to use T0pp but T5 is close enough)
- T5, BART, Marian, Pegasus, EncoderDecoder: @patrickvonplaten

Library:
- Deepspeed: @stas00
- Text generation: @patrickvonplaten @narsil

## Information

Model I am using (Bert, XLNet ...): T0pp / T0_3B

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [X] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [X] my own task or dataset: (give details below)

## To reproduce
I want to load T0pp across 2 24GB GPUs and only run inference. I know Deepspeed wit zeRO stage 3 is the way to go for this from reading documentation. I am following the HuggingFace example [here](https://huggingface.co/docs/transformers/v4.16.1/en/main_classes/deepspeed#deepspeed-non-trainer-integration) to use Deepspeed without a `Trainer` object.

The error I get is 
```
[2022-01-28 18:36:41,193] [INFO] [partition_parameters.py:456:__exit__] finished initializing model with 2.85B parameters
Traceback (most recent call last):
  File "multi_gpu_T0pp.py", line 26, in <module>
    engine = deepspeed.initialize(model=model, config_params=ds_config)
AttributeError: module 'transformers.deepspeed' has no attribute 'initialize'
```

My code:

Run with `CUDA_VISIBLE_DEVICES="0,1" deepspeed <script.py>`

```python
"""
Example code to load a PyTorch model across GPUs
"""
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from transformers.deepspeed import HfDeepSpeedConfig
from transformers import deepspeed
import pandas as pd
import torch
import pdb
import os

seed = 42
torch.manual_seed(seed)

ds_config = {
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "zero_optimization": {
        "stage": 3,
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_fp16_weights_on_model_save": true
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": 0,
    "steps_per_print": 2000,
    "train_batch_size": 2,
    "train_micro_batch_size_per_gpu": 1,
    "wall_clock_breakdown": false
}

if __name__ == "__main__":
    # must run before instantiating the model
    # ds_config is deepspeed config object or path to the file
    dschf = HfDeepSpeedConfig(ds_config)  # keep this object alive

    model_name = "bigscience/T0_3B"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

    engine = deepspeed.initialize(model=model, config_params=ds_config)

    inputs = tokenizer.encode(
        "Is this review positive or negative? Review: this is the best cast iron skillet you will ever buy",
        return_tensors="pt")
    outputs = model.generate(inputs)
    print(tokenizer.decode(outputs[0]))
```

## Expected behavior
T0pp (or T0_3B) to load across 2 GPUs, generate an answer, and then quit.
