michaelfeil / michaelfeil/infinity
Permit loading of models at different precision at load time for sentence_transformers
- 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