michaelfeil / michaelfeil/infinity

float16 and other optimizations help?

Open
#159 6 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
2.9k
Forks
206
PR merge metrics
No merged PRs in 30d

Description

I built the following script based on reviewing yours...thanks BTW. It only implements float16 but I'd love to incorporate other quantizations and/or better transformer or whatever else. Hope what I spend a good day on helps you in your endeavors:

```
import torch
from sentence_transformers import SentenceTransformer, util
import numpy as np
import time

class CustomSentenceTransformer(SentenceTransformer):
def __init__(self, model_name_or_path="sentence-transformers/all-mpnet-base-v2", device='cuda', dtype=torch.float16):
# Call the superclass initializer directly with the device. The superclass should handle device placement.
super().__init__(model_name_or_path, device=device)

if torch.cuda.is_available() and device == 'cuda' and dtype == torch.float16:
self.half()
print("Model converted to float16 precision.")

def encode(self, sentences, batch_size=32, show_progress_bar=False, convert_to_numpy=True, normalize_embeddings=True):
self.eval()

all_embeddings = []
for start_index in range(0, len(sentences), batch_size):
sentences_batch = sentences[start_index:start_index+batch_size]
features = self.tokenize(sentences_batch)
features = util.batch_to_device(features, self.device)

with torch.no_grad():
out_features = self.forward(features)
embeddings = out_features['sentence_embedding']

if normalize_embeddings:
embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)

all_embeddings.append(embeddings)

all_embeddings = torch.cat(all_embeddings, 0)

if convert_to_numpy:
all_embeddings = all_embeddings.cpu().numpy()

return all_embeddings
```

By the way, I'm noticing 1/2 compute time but the VRAM usage isn't going down at all. The model I'm using was originally created in float16 so...What I'm "guessing" is that even though it's originally float16, prior to using this script, it runs in float32 because dtype isn't set to float16? Guessing that's why you named your class "patched?"

At any rate, would be curious to know your expert opinion. Here's most of the function I'm using with it...

```
def generate_embeddings(text_chunks):
model_name = r"PATH HERE"
model_kwargs = {'device': 'cuda'}
encode_kwargs = {'normalize_embeddings': True}

model = CustomSentenceTransformer(model_name, **model_kwargs)

start_time = time.time()

embeddings = model.encode(text_chunks, **encode_kwargs)

end_time = time.time()
elapsed_time = end_time - start_time

print(f"Embedding generation took {elapsed_time:.2f} seconds.")
```

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reviewing the CustomSentenceTransformer class and the generate_embeddings function shown in the issue. Compare the float16 behavior and VRAM usage, then investigate which additional quantization or transformer optimizations are appropriate. Done would require a clearly scoped optimization with measured compute time and memory results.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
20/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.