running spacy for predicting ner on mulitple GPUs

Question

I am using spacy for predicting ner tags using gpu. I have a bigger machine with 8 GPUs, I want to use all those GPUs.

one way to use all gpus are running 8 different script on each gpu and then passing text to each script using kafka queue.

Is there any other way through which we can use single script to use all the gpus for predicting ner.

Answer 1

Have you tried running Spacy with MPI? I am experimenting the following code myself, so please let me know if it works!

from mpi4py import MPI
import cupy

comm = MPI.COMM_WORLD
rank = comm.Get_rank()

if rank == 0:
    data = ["His friend Nicolas J. Smith is here with Bart Simpon and Fred."*100]
else:
    data = None

unit = comm.scatter(data, root=0)

with cupy.cuda.Device(rank):
    import spacy
    from thinc.api import set_gpu_allocator, require_gpu
    set_gpu_allocator("pytorch")
    require_gpu(rank)
    nlp = spacy.load('en_core_web_lg')
    nlp.add_pipe("merge_entities")
    tmp_list = []
    for doc in nlp.pipe(unit):
        res = " ".join([t.text if not t.ent_type_ else t.ent_type_ for t in doc])
        tmp_list.append(res)

result = comm.gather(tmp_list, root=0)

if comm.rank == 0:
    print (result)
else:
    result = None