How to Use Transformers pipeline with multiple GPUs
Author: vikramtharakan
Pull request: no

I'm using an tweaked version of the `uer/roberta-base-chinese-extractive-qa` model. While I know how to train using multiple GPUs, it is not clear how to use multiple GPUs when coming to this stage. Essentially this is what I have:

```
from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

model = AutoModelForQuestionAnswering.from_pretrained('uer/roberta-base-chinese-extractive-qa')
tokenizer = AutoTokenizer.from_pretrained('uer/roberta-base-chinese-extractive-qa')
QA = pipeline('question-answering', model=model, tokenizer=tokenizer, device=0)
```

Except the model I'm using is a tweaked version of the base model shown above. I'm trying to process a large number of documents and at the rate it's going at it would take on the order of weeks. Is there anyway to set the pipeline to use multiple gpus? So I could use something like

```
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
device = torch.device("cuda")
```
And then use both devices?

Also looks like someone else had the same question: https://stackoverflow.com/questions/64193049/how-to-use-transformers-pipeline-with-multi-gpu


Comments:
cc @Narsil 

Hi @vikramtharakan,

Currently no, it's not possible in the pipeline to do that.
We not ruling out putting it in at a later stage, but it's probably a very involved process, because there are many ways someone could want to use multiple GPUs for inference.
- If the model fits a single GPU, then get parallel processes, 1 on all GPUs and run inference on those
- If the model doesn't fit a single GPU, then there are multiple options too, involving deepspeed or JaX or TF tools to handle model parallelism, or data parallelism or all of the, above.


If your model fits a single GPU, then you could easily have something like:

`myfile.py`
```python
import os

START=int(os.getenv("START"))
STOP=int(os.getenv("STOP"))
DEVICE=os.getenv("DEVICE")

for result in pipe(dataset[START:STOP], device=DEVICE):
    print(result)
   .. # logic
``` 

Then start a single process on each GPU with something like
```
START=0 STOP=100 DEVICE=0 python myfile.py &
START=100 STOP=200 DEVICE=1 python myfile.py &
```

This would launch a single process per GPU, with controllable access to the dataset and the device.

Would that sort of approach work for you ?


Note: In order to feed the GPU as fast as possible, the pipeline uses a DataLoader which has the option `num_workers`. A good default would be to set it to `num_workers = num_cpus (logical + physical) / num_gpus` so that you don't overdo parallelism. 

Note2: As in all things considering performance, measure and keep measuring there are many things that could go wrong when attempting to maximise performance hardware, sometimes unexpectedly

@Narsil sorry for the late reply; I understand what you're saying and I think I will have to implement something to this effect in practice but what you suggested won't quite work for me. The model does fit a single GPU but ideally I'd like data to be able to flow through it continuously. So we'd essentially have one pipeline set up per GPU that each runs one process, and the data can flow through with each context being randomly assigned to one of these pipes using something like python's multiprocessing tool, and then aggregate all the data at the end. Do you have any tips on how to implement this? I'm struggling to find the best way to implement this

> Do you have any tips on how to implement this?

Could you have a queue instead maybe of choosing the offsets in the datasets ?
Both processes could be iteratively fetching from the queue (that can be fed in random order or not).
Hashing the id `if hash(id) % num_processes == my_process_id: process else: skip`can be used if you want to fetch in random order without a queue but making sure you're not duplicating the work
Many programs can be used to implement a queue simply, rabbitmq, mongo, maybe the pure Python `multiprocessing.Queue`. 

That's the first idea to come to mind, but I would have many questions to assess the best solution.

Do you want to be able to steal work from others because you are sensitive to latency ?
Do you not care and have a simple round-robin ?
Why is the randomness important ?
Is you data very large and not easy to pass around in queue system X, do you need shared memory ?
Are you passing memory efficiently down on each GPU ?
Is the data on disk or RAM ?

This is again beyond of the scope of `pipelines` since there are many solutions to this problem where there's no obvious better solution and it all depends on what you want to achieve and what you are your constraints.

@Narsil Gotcha, thank you for the quick reply! I'll look into that

I was also wondering if there's any documentation on all the steps going on under the hood when using these pipelines? Because if it's really just tokenizing the data and then sending it through the model, I feel like I could also just load the model separately using pytorch and then load the model to multiple devices using either pytorch's `DataParallel` or `DistributedDataParallel`. The tokenization could be done outside of that as well, as I imagine it doesn't need the GPU as much as the model will for making predictions

Hi @vikramtharakan ,

Pipeline are intended for users not knowing machine learning to provide easy to use API.
It should work with the sanest possible defaults, and have a very standard preprocessing/postprocessing.
It is not meant to be the best possible production ready inference tool on all hardware (This is too vast a problem to be solved by a single solution).

It does use DataLoader when available, it makes the batching works even on non-obvious inputs (like question answering with long contexts or speech recognition on long audio files).
Currently there's no recap of all the features they provide, that's a good thing to add this.

But if you want fine grained over your inference loops, using the bare tools FeatureExtractor/Tokenizer/AutoModel is definitely the easiest way to take control over everything. You will just have to write a lot of boilerplate code again.
You can also subclass your pipeline class, but that won't make multi GPU easy. At the core the pipeline assume s1 model only.

Again personally, I would first try 1 pipeline per process as that seems pretty easy way to do things and this is how we use it internally for the API, but your mileage may vary and `DataParallel` is also a good choice.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

For compatible models, [parallelformers](https://github.com/tunib-ai/parallelformers) models work as a drop-in replacement. 

> Hi @vikramtharakan ,
> 
> Pipeline are intended for users not knowing machine learning to provide easy to use API. It should work with the sanest possible defaults, and have a very standard preprocessing/postprocessing. It is not meant to be the best possible production ready inference tool on all hardware (This is too vast a problem to be solved by a single solution).
> 
> It does use DataLoader when available, it makes the batching works even on non-obvious inputs (like question answering with long contexts or speech recognition on long audio files). Currently there's no recap of all the features they provide, that's a good thing to add this.
> 
> But if you want fine grained over your inference loops, using the bare tools FeatureExtractor/Tokenizer/AutoModel is definitely the easiest way to take control over everything. You will just have to write a lot of boilerplate code again. You can also subclass your pipeline class, but that won't make multi GPU easy. At the core the pipeline assume s1 model only.
> 
> Again personally, I would first try 1 pipeline per process as that seems pretty easy way to do things and this is how we use it internally for the API, but your mileage may vary and `DataParallel` is also a good choice.

Hi, @Narsil ,

I have the same problem and I agree with your solution. But I had some trouble implementing it. I used spawn of torch, tried to implement shared variables with Manager, nothing worked. Do you have any good ways to implement it?

@hasakikiki 

It depends really as I said.
What model are you trying to load, on what hardware ?

Do you need to split the model itself on various GPUs (like `bloom` for instance)  ?
Do you even need multiple nodes for this ?

Do you have 1 node with several GPUs and just want to be able to use them in parallel ? (Then using `pipeline(.., device=n)` should suffice for instance.

I unfortunately don't really know what is torch `Manager` and shared variables.

@Narsil 

I'm sorry my description is too vague.

Yes, I have a machine with 8 GPUs. And I don't need to split the model. I want to use them in parallel during inference. But in the `pipeline` API documentation, it seems that only one GPU can be used, which is the `device=n` you provided.

> Yes, I have a machine with 8 GPUs. And I don't need to split the model. I want to use them in parallel during inference. But in the pipeline API documentation, it seems that only one GPU can be used, which is the device=n you provided.

Then just start 8 processes each with a pipeline attributed on a single GPU, and distribute the work through a queue of some kind like redis/rabbitmq or any other similar mecanism ?

@Narsil I am trying to run the feature extraction pipeline on multiple GPU via composing the code using snippets you share:
```python
import os
import numpy as np
import hickle as hkl
import torch
from torch.utils.data import Dataset
from tqdm import tqdm
tqdm.pandas()

from transformers import AutoModel, pipeline
from transformers import AutoTokenizer
from torch.multiprocessing import Pool, Process, set_start_method, get_context

set_start_method("spawn", force=True)

pipe = None
model_name="anferico/bert-for-patents"

class NoDaemonProcess(Process):
    @property
    def daemon(self):
        return False

    @daemon.setter
    def daemon(self, value):
        pass

class NoDaemonContext(type(get_context("fork"))):
    Process = NoDaemonProcess


class ListDataset(Dataset):
    def __init__(self, original_list):
        self.original_list = original_list
    def __len__(self):
        return len(self.original_list)
    def __getitem__(self, i):
        return self.original_list[i]


def get_pipe(device_index):
    # This will load the pipeline on demand on the current PROCESS/THREAD.
    # And load it only once.
    global pipe
    if pipe is None:
        print(f"Creating pipeline for device: {device_index}")
        tokenizer = AutoTokenizer.from_pretrained(model_name, do_lower_case=True, model_max_length=512, truncation=True,
                                                  padding=True, pad_to_max_length=True)
        device = torch.device(f'cuda:{device_index}' if torch.cuda.is_available() else 'cpu')
        model = AutoModel.from_pretrained(model_name).to(device)
        pipe = pipeline('feature-extraction', model=model, tokenizer=tokenizer, max_length=512, truncation=True,
                        padding=True, pad_to_max_length=True, device=device, framework="pt", batch_size=16)
    return pipe


def extract_chunk_embedding(dataset, device_index):
    extractor = get_pipe(device_index)
    print(f"Getting embeddings for device: {device_index}, chunked data length: {len(dataset)}")
    return [np.squeeze(embedding)[0] for embedding in
            tqdm(extractor(dataset, max_length=512, truncation=True, num_workers=8), total=len(dataset))]


def extract_embeddings(df, col, prefix, mode="cls", overwrite=False, num_gpu=2):
    embedding_file = f'transformer_embeddings_cls_{prefix}.hkl'
    if not overwrite:
        if mode == "cls" and os.path.isfile(embedding_file):
            try:
                embeddings_dict = hkl.load(embedding_file)
                print(f"Loading from the disk: {embedding_file}...")
                return df['text_hash'].map(embeddings_dict).to_list()
            except Exception as err:
                print(f"Error reading file: {embedding_file} -> {err}")
                pass

    dataset = ListDataset(df[col].tolist())
    if mode == "cls":
        print(f"Creating embeddings: {embedding_file}...")
        num_gpu = min(num_gpu, torch.cuda.device_count())
        mp_context = NoDaemonContext()
        multi_pool = mp_context.Pool(processes=num_gpu)
        data_indices = tuple(np.ceil(np.linspace(start=0, stop=len(dataset), num=num_gpu + 1)).astype(np.int32))
        arg_list = []
        for device_index, (start, stop) in enumerate(zip(data_indices, data_indices[1:])):
            chunked_data = dataset[start: stop]
            arg_list.append((chunked_data, device_index))

        print("Starting multi-processing...")
        res = multi_pool.starmap(extract_chunk_embedding, arg_list)
        multi_pool.close()
        multi_pool.join()

        return res
```
So the input is a data frame with a column name in which the column has text data and I am trying to extract features. Unfortunately, it is hanging on the multiprocessing step, and I could not find why it is the case, thank you for your answer.

Could you try using:

```python
from transformers import pipeline, AutoTokenizer, AutoModel
import torch
from torch.utils.data import Dataset
from multiprocessing import Pool


MODEL_NAME = "anferico/bert-for-patents"


def get_pipeline(device_index):
    print(f"Creating pipeline for device: {device_index}")
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_NAME, do_lower_case=True, model_max_length=512, truncation=True, padding=True, pad_to_max_length=True
    )
    device = torch.device(f"cuda:{device_index}" if torch.cuda.is_available() else "cpu")
    model = AutoModel.from_pretrained(MODEL_NAME).to(device)
    pipe = pipeline(
        "feature-extraction",
        model=model,
        tokenizer=tokenizer,
        max_length=512,
        truncation=True,
        padding=True,
        pad_to_max_length=True,
        device=device,
        framework="pt",
        batch_size=16,
    )
    return pipe


dummy = ["This is a test"] * 100


class MyDataset(Dataset):
    def __init__(self, device_index, world_size):
        super().__init__()
        self.device_index = device_index
        self.data = dummy[device_index::world_size]
        print(self.data)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, i):
        return self.data[i]


def run(arg):
    device_index, world_size = arg
    pipe = get_pipeline(device_index)

    data = MyDataset(device_index, world_size)

    for out in pipe(data):
        pass
        # Do something here, directly on the process
        # Out is what you want on a single item


def main():
    num_gpus = 8
    with Pool(num_gpus) as p:
        p.map(run, [(i, num_gpus) for i in range(num_gpus)])


if __name__ == "__main__":
    main()
```

I escaped the issues with pandas and your dataset which I don't know, but the idea is to handle the loading directly in the dataset, and to create the pipelines only once, and directly within the subprocess.

Also I recommend that you treat directly results as you receive them so that you don't have to save n tensors and grow your memory usage indefinitely.

> You are using

@Narsil I do not get it.

> @Narsil I do not get it.

I started the comment, and misfired send as obviously it wasn't ready.

Is there a fix for this ?
I want to run `GPT-NeoX-20b` on multiple gpus as a single one is not possible

As I said earlier, it's a complex problem and there's no one solution to making things Parallel.

**Please be considerate or your problem space and think about how you going about things. I can't stress that enough, every use has different potential solutions**
There is no 1-click solution for all options.

That being said, since this first comment was published the libs did make progress and you can now use `accelerate` and use

```python
pipe = pipeline(...., device_map="auto")
```
Which will attempt to put the model on multiple GPUs.

This method works by putting the splitting the model across layers, and putting layer 0, 1, 2 on GPU 0 and layer 3, 4, 5 on GPU1 (it's an example).

This is extremely simple, but it's possibly the worst solution if you care about latency. If you care about throughput, then you are not optimal either since you could make things multiprocessed to be able to process different parts of data on each GPU separately.

But if you just want to take the model out for a spin, this is the easiest solution.

I want to use GPT-NeoX on AWS EC2 g4dn.12xlarge.

GPUs: 4
GPU model: NVIDIA Tesla M60
GPU memory: 32 GiB

Is there any better option than pipe = pipeline(...., device_map="auto")?

What option will work?
pipe = pipeline(model="EleutherAI/gpt-neox-20b", device_map="auto")
or
pipe = pipeline(model="EleutherAI/gpt-neox-20b", device=0, device_map="auto")

You cannot specify both `device=0` (which means put everything on the first GP) and `device_map="auto"` ( which will try to put everything on each GPU as it sees fit, it will even use CPU and disk if needed). 

So unfortunately no there is no better option than

```python
pipe = pipeline(...., device_map="auto")
```
as least as a one line thing.

For more complete solutions, there's probably a lot of blogpost and tutorials out there.

For a webserver trying to optimize latency you can check out : https://github.com/huggingface/text-generation-inference/ which is a repo we're using for bloom. NB: Do not expect the same level of support as in core transformers since this is meant as an internal tool (we're just publishing it so others can see/improve and use it).
It does quite a few things, by batching queries dynamically, using custom kernels (not available for neox) and using Tensor Parallelism instead of Pipeline Parallelism (what `accelerate` does).  

@Narsil How can we use multiple GPUs in pipeline for inference? When I set `device_map="auto"`, it shows
```
Hardware accelerator e.g. GPU is available in the environment, but no `device` argument is passed to the `Pipeline` object. Model will be on CPU.
```

However, when I set `device="cuda", device_map="auto"`, it shows
```
Both `device` and `device_map` are specified. `device` will override `device_map`. You will most likely encounter unexpected behavior. Please remove `device` and keep `device_map`.
```
and only uses one GPU.

cc @Rocketknight1 as well

@yujianll that's an interesting issue - can you paste me some code to reproduce it? I don't get that warning when using `device_map="auto"`.