Limiting Memory Consumption

[thatcort]

Hi, we're trying to use Spacy for sentence segmentation (and ideally NER at the same time), but running into GPU memory issues. It seems that the amount of memory used increases with the number of sentences in the documents. This is causing OOM errors and other tasks on the GPU to fail or not run. We're trying to parameterize this to work in a Ray cluster where we allocate a certain amount of memory to the sentence segmentation task in advance (we can't know how many sentences there will be before running!).

Is there a way to limit the amount of memory that Spacy will use on the GPU?

Thank you!

[adrianeboyd]

Depending on the type of model, spacy is using either cupy or pytorch to allocate GPU memory, so you can control the memory usage through cupy or pytorch options:

https://docs.cupy.dev/en/stable/user_guide/memory.html

https://pytorch.org/docs/2.1/generated/torch.cuda.set_per_process_memory_fraction.html#torch.cuda.set_per_process_memory_fraction

By default, the sm/md/lg models use cupy (since torch is not a requirement) and the trf models use pytorch. If you'd prefer to use pytorch options across the board, I think you should be able to switch to pytorch through thinc utils (but I'm not entirely sure what the overhead is vs. cupy-only allocation):

https://thinc.ai/docs/api-util#use_pytorch_for_gpu_memory