How to use transformers pipeline with multi-gpu?

Question

ner_model = pipeline('ner', model=model, tokenizer=tokenizer, device=0, grouped_entities=True)

the device indicated pipeline to use no_gpu=0(only using GPU), please show me how to use multi-gpu.

Answer 1

There is an argument called device_map for the pipelines in the transformers lib; see here. It comes from the accelerate module; see here. You can specify a custom model dispatch, but you can also have it inferred automatically with device_map=" auto". Eventually, you might need additional configuration for the tokenizer, but it should look like this:

ner_model = pipeline('ner', model=model, tokenizer=tokenizer, device_map="auto", grouped_entities=True)