Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Updated to allow the selection of GPU for embedding where there is mo… #1734

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Updated to allow the selection of GPU for embedding where there is mo… #1734

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
207531d
Updated to allow the selection of GPU for embedding where there is mo…
sgresham Mar 15, 2024
16ee058
gpufix partial
sgresham Mar 20, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 23 additions & 1 deletion private_gpt/components/embedding/embedding_component.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
from private_gpt.settings.settings import Settings

logger = logging.getLogger(__name__)

import torch
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd move this to the try block within "huggingface" case. There is no "torch" general dependency declared in pyproject.toml, so this could break the whole execution for people not using huggingface. Actually, we may need to add torch to embeddings-huggingface = ["llama-index-embeddings-huggingface"] as

# Optional Huggingface related dependency
torch = {version = "^2.2.1", optional = true}

embeddings-huggingface = ["torch", "llama-index-embeddings-huggingface"]

in pyproject.toml.

I think huggingface package from llamaindex already depends on torch, but given we are now importing it explicitly we should also depende on it.


@singleton
class EmbeddingComponent:
Expand All @@ -28,9 +28,31 @@ def __init__(self, settings: Settings) -> None:
"Local dependencies not found, install with `poetry install --extras embeddings-huggingface`"
) from e

# Get the number of available GPUs
num_gpus = torch.cuda.device_count()
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Adding code to the codebase just to print information is not a good practive. I'd remove this. whole block of prints.


if num_gpus > 0:
print("Available CUDA devices:")
for i in range(num_gpus):
print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
else:
print("No CUDA devices available. Switching to CPU.")

# Check if CUDA is available
if torch.cuda.is_available():
# If settings.embedding.gpu is specified, use that GPU index
if hasattr(settings, 'huggingface') and hasattr(settings.huggingface, 'gpu_type'):
device = torch.device(f"{settings.huggingface.gpu_type}:{settings.huggingface.gpu_number}")
else:
device = torch.device('cuda:0')
else:
# If CUDA is not available, use CPU
device = torch.device("cpu")
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What happens with laptops using a GPU that is not Nvidia based? For example Mac book running Metal GPU? Will this make embedding slower forcing them to go to CPU?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This logic looks similar this: llama_index.core.utils.infer_torch_device which handles Metal (mps).

print("Embedding Device: ",device)
self.embedding_model = HuggingFaceEmbedding(
model_name=settings.huggingface.embedding_hf_model_name,
cache_folder=str(models_cache_path),
device=device
)
case "sagemaker":
try:
Expand Down
9 changes: 8 additions & 1 deletion private_gpt/settings/settings.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Literal
from typing import Literal, Optional

from pydantic import BaseModel, Field

Expand Down Expand Up @@ -145,6 +145,13 @@ class HuggingFaceSettings(BaseModel):
embedding_hf_model_name: str = Field(
description="Name of the HuggingFace model to use for embeddings"
)
gpu_type: Optional[Literal["cuda","cpu"]] = Field(
description="GPU typedevice for embedding, can be 'cuda' or cpu"
)
gpu_number: int = Field(
0,
description="GPU device number for embedding, will be presented to torch like 'cuda:x'"
)


class EmbeddingSettings(BaseModel):
Expand Down
4 changes: 3 additions & 1 deletion settings.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,9 +54,11 @@ embedding:
# Should be matching the value above in most cases
mode: huggingface
ingest_mode: simple

huggingface:
embedding_hf_model_name: BAAI/bge-small-en-v1.5
gpu_type: cuda #GPU typedevice for embedding, can be 'cuda', rocm or cpu". defaults to cuda[0], or cpu if cuda not available
gpu_number: 1 #Directly select a device, normally 0 if only a single GPU

vectorstore:
database: qdrant
Expand Down