michaelfeil / michaelfeil/infinity

Permit loading of models at different precision at load time for sentence_transformers

Open
#331 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
2.9k
Forks
206
PR merge metrics
No merged PRs in 30d

Description

### Feature request

Pass in `torch_dtype` in model_kwargs, as supported by sentence_transformers when specifying dtype in the infinity_emb v2 cli when InferenceEngine type is torch.

This would allow the loading of the Transformer model at a lower precision at load time instead of post-loading, which could cause an OOM error.

Post loading quantization, e.g. self.half(), would still be needed since it appears the non-transformer Pytorch modules in some models are still in fp32 and would cause issues downstream where matrix computations later would fail due to mixed type computation.

### Motivation

While the current code quantizes the precision post loading a model, the issue is that if we're loading full 32-bit float models against a GPU that has a small amount of memory, it could fail to load since the total size of the model would exceed the GPUs maximum memory limit. A use case of this would be seen in small multi-instance GPUs e.g. deploying a 10Gb GPU instance in an NVIDIA A100 in MIG mode.

By specifying the precision at load time, we would be able to load a model without OOM errors and successfully use the model.

### Your contribution

Yes. Happy to submit a PR.

Current tests on my end utilize the following code modifications in the init function of the SentenceTransformerPatched class.

```
if engine_args.engine == InferenceEngine.torch and \
engine_args.device in [Device.auto, Device.cuda] and \
engine_args.dtype == Dtype.float16:
model_kwargs["torch_dtype"] = torch.float16
```
but would need further work for better support/handling of other types.
Some thoughts on what can be done for each dtype per the torch attributes here: https://pytorch.org/docs/stable/tensor_attributes.html

auto -> (leave blank)
float32 -> torch.float
float16 -> torch.half
float8 -> (leave blank) ?
int8 -> (leave blank) ?

The 8-bit options were left blank since I noticed that quantization is performed later for these types - unsure if there is a better approach here.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start in the SentenceTransformerPatched __init__ method, where the issue says model_kwargs and engine_args are handled. Trace how dtype values reach sentence_transformers and how later quantization is applied; done means supported load-time dtypes avoid the described GPU OOM while preserving compatible post-loading behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.