[Question]: embedding model support multi-gpu

[RuiqingGuo]

Question Validation

• [X] I have searched both the documentation and discord for an answer.

Question

I want to use the semantic splitter from llamaindex for document segmentation. Is there any support for embedding model parallelism across multiple gpus? Here is the code I used with one GPU(A40) and the error message I encountered.

dir_path = "/data/test-splitter"
documents = SimpleDirectoryReader(dir_path).load_data()
embed_dir = '/models/embedding/bge-m3'
embed_model = HuggingFaceEmbedding(model_name=embed_dir, device="coda:0",
                                   embed_batch_size=1)
splitter = SemanticSplitterNodeParser(
    buffer_size=1, breakpoint_percentile_threshold=50, embed_model=embed_model
)
nodes = splitter.get_nodes_from_documents(documents)
for node in nodes:
    print('-' * 100)
    print(node.get_content())

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.46 GiB. GPU 0 has a total capacity of 44.56 GiB of which 1.17 GiB is free. Process 2617384 has 43.38 GiB memory in use. Of the allocated memory 42.33 GiB is allocated by PyTorch, and 760.95 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

[dosubot[bot]]

Hello, @RuiqingGuo! I'm here to help you with any bugs, questions, or contributions you may have. Let's tackle this issue together!

Yes, there is support for embedding model parallelism across multiple GPUs when using the semantic splitter from LlamaIndex for document segmentation. The provided code includes configurations for tensor parallelism (tp_size) and pipeline parallelism (pp_size), which are used to distribute the model across multiple GPUs.

tp_size = config["builder_config"]["tensor_parallel"]
pp_size = config["builder_config"]["pipeline_parallel"]
world_size = tp_size * pp_size
assert (
    world_size == tensorrt_llm.mpi_world_size()
), f"Engine world size ({world_size}) != Runtime world size ({tensorrt_llm.mpi_world_size()})"

However, it is noted that the Python runtime does not support pipeline parallelism:

assert (
    pp_size == 1
), "Python runtime does not support pipeline parallelism"

This means that while tensor parallelism is supported, pipeline parallelism is not supported in the Python runtime ^[1].

Additionally, the HuggingFaceEmbedding model does not include explicit handling for multi-GPU setups. The _device attribute is set to a single device, and the model is moved to this single device:

self._device = device or infer_torch_device()
self._model = model.to(self._device)

This indicates that the current implementation does not support distributing the model across multiple GPUs ^[2].

For better multi-GPU support, consider using alternative embedding models such as OptimumEmbedding, OpenVINOEmbedding, TogetherEmbedding, IpexLLMEmbedding, OctoAIEmbedding, VoyageEmbedding, or GradientEmbedding ^[3][4].

To continue talking to Dosu, mention @dosu.

[thorory]

i meet the same problem with a large pdf(260 pages).

when I use SimpleDirectoryReader, it returns documents for every page. And then SemanticSplitterNodeParser can embed the text. However, when I use my own OCR and put all the content in one string, the SemanticSplitterNodeParser get OOM