如何在多 GPU 上使用 Transformer 管道？

Question

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

设备指示管道使用 no_gpu=0（仅使用 GPU），请告诉我如何使用多 GPU 。

Answer 1

有一个参数要求libdevice_map中的管道transformers；看这里。它来自accelerate模块；看这里。您可以指定自定义模型调度，但也可以使用device_map=" auto". 最终，您可能需要对标记生成器进行额外的配置，但它应该如下所示：

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