Running a SOTA 7B Parameter Embedding Model on a Single GPU | by Szymon Palucha

Contents

Set Up How to Run — In Detail Annotated Example How to Run — TLDR

Set Up

The model that we will experiment with is the Alibaba-NLP/gte-Qwen2-7B-instruct from Transformers. The model card is here.

To perform this experiment, I have used Python 3.10.8 and installed the following packages:

torch==2.3.0
transformers==4.41.2
xformers==0.0.26.post1
flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.5.8/flash_attn-2.5.8+cu122torch2.3cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
accelerate==0.31.0

I ran into some difficulty in installing flash-attn required to run this model and so had to install the specific version listed above. If anyone has a better workaround please let me know!

The Amazon SageMaker instance I used for this experiment is the ml.g5.2xlarge. It has a 24GB NVIDIA A10G GPU and 32GB of CPU memory and it costs $1.69/hour. The below screenshot from AWS shows all the details of the instance

SageMaker g5 instance types from AWS docs.

Actually to be precise if you run nvidia-smi you will see that the instance only has 23GB of GPU memory which is slightly less than advertised. The CUDA version on this GPU is 12.2.

How to Run — In Detail

If you look at the model card, one of the suggested ways to use this model is via the sentence-transformers library as show below

from sentence_transformers import SentenceTransformer# This will not run on our 24GB GPU!
model = SentenceTransformer("Alibaba-NLP/gte-Qwen2-7B-instruct", trust_remote_code=True)
embeddings = model.encode(list_of_examples)

Sentence-transformers is an extension of the Transformers package for computing embeddings and is very useful as you can get things working with two lines of code. The downside is that you have less control on how to load the model as it hides away tokenisation and pooling details. The above code will not run on our GPU instance because it attempts to load the model in full float32 precision which would take 28GB of memory. When the sentence transformer model is initialised it checks for available devices (cuda for GPU) and automatically shifts the Pytorch model onto the device. As a result it gets stuck after loading 5/7ths of the model and crashes.

Instead we need to be able to load the model in float16 precision before we move it onto the GPU. As such we need to use the lower level Transformers library. (I am not sure of a way to do it with sentence-transformers but let me know if one exists!) We do this as follows

import transformers
import torchmodel_path = "Alibaba-NLP/gte-Qwen2-7B-instruct"
model = transformers.AutoModel.from_pretrained(model_path, trust_remote_code=True, torch_dtype=torch.float16).to("cuda")

With the torch_dtype parameter we specify that the model should be loaded in float16 precision straight away, thus only requiring 14GB of memory. We then need to move the model onto the GPU device which is achieved with the to method. Using the above code, the model takes almost 2min to load!

Since we are using transformers we need to separately load the tokeniser to tokenise the input texts as follows:

tokenizer = transformers.AutoTokenizer.from_pretrained(model_path)

The next step is to tokenise the input texts which is done as follows:

texts = ["example text 1", "example text 2 of different length"]
max_length = 32768
batch_dict = tokenizer(texts, max_length=max_length, padding=True, truncation=True, return_tensors="pt").to(DEVICE)

The maximum length of the Qwen2 model is 32678, however as we will see later we are unable to run it with such a long sequence on our 24GB GPU due to the additional memory requirements. I would recommend reducing this to no more than 24,000 to avoid out of memory errors. Padding ensures that all the inputs in the batch have the same length whilst truncation ensures that any inputs longer than the maximum length will be truncated. For more information please see the docs. Finally, we ensure that we return PyTorch tensors (default would be lists instead) and move these tensors onto the GPU to be available to pass to the model.

The next step is to pass the inputs through our model and perform pooling. This is done as follows

with torch.no_grad():
outputs = model(**batch_dict)
embeddings = last_token_pool(outputs.last_hidden_state, batch_dict["attention_mask"])

with the last_token_pool which looks as follows:

def last_token_pool(last_hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
# checks whether there is any padding (where attention mask = 0 for a given text)
no_padding = attention_mask[:, -1].sum() == attention_mask.shape[0]
# if no padding - only would happen if batch size of 1 or all sequnces have the same length, then take the last tokens as the embeddings
if no_padding:
return last_hidden_states[:, -1]
# otherwise use the last non padding token for each text in the batch
sequence_lengths = attention_mask.sum(dim=1) - 1
batch_size = last_hidden_states.shape[0]
return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengthsLet’s break down what happened in the above code snippets!

The torch.no_grad() context manager is used to disable gradient calculation, since we are not training the model and hence to speed up the inference.
We then pass the tokenised inputs into the transformer model.
We retrieve the outputs from the last layer of the model with the last_hidden_state attribute. This is a tensor of shape (batch_size, max_sequence_length, embedding dimension). Essentially for each example in the batch the transformer outputs embeddings for all the tokens in the sequence.
We now need some way of combining all the token embeddings into a single embedding to represent the input text. This is called pooling and it is done in the same way as during training of the model.
In older BERT based models the first token was typically used (which represented the special classification [CLS] token). However, the Qwen2 model is LLM-based, i.e. transformer decoder based. In the decoder, the tokens are generated auto regressively (one after another) and so the last token contains all the information encoded about the sentence.
The goal of the last_token_pool function is to therefore select the embedding of the last generated token (which was not the padding token) for each example in the batch.
It uses the attention_mask which tells the model which of the tokens are padding tokens for each example in the batch (see the docs).

Annotated Example

Let’s look at an example to understand it in a bit more detail. Let’s say we want to embed two examples in a single batch:

texts = ["example text 1", "example text 2 of different length"]

The outputs of the tokeniser (the batch_dict ) will look as follows:

>>> batch_dict
{'input_ids': tensor([[  8687,   1467,    220,     16, 151643, 151643, 151643],
[  8687,   1467,    220,     17,    315,   2155,   3084]],
device='cuda:0'), 'attention_mask': tensor([[1, 1, 1, 1, 0, 0, 0],
[1, 1, 1, 1, 1, 1, 1]], device='cuda:0')}

From this you can see that the first sentence gets split into four tokens (8687, 1467, 220, 16), while the second sentence get split into seven tokens. As a result, the first sentence is padded (with three padding tokens with id 151643) up to length seven — the maximum in the batch. The attention mask reflects this — it has three zeros for the first example corresponding to the location of the padding tokens. Both the tensors have the same size

>>> batch_dict.input_ids.shape
torch.Size([2, 7])
>>> batch_dict.attention_mask.shape
torch.Size([2, 7])

Now passing the batch_dict through the model we can retrieve the models last hidden state of shape:

>>> outputs.last_hidden_state.shape
torch.Size([2, 7, 3584])

We can see that this is of shape (batch_size, max_sequence_length, embedding dimension). Qwen2 has an embedding dimension of 3584!

Now we are in the last_token_pool function. The first line checks if padding exists, it does it by summing the last “column” of the attention_mask and comparing it to the batch_size (given by attention_mask.shape[0]. This will only result in true if there exists a 1 in all of the attention mask, i.e. if all the examples are the same length or if we only have one example.

>>> attention_mask.shape[0]
2
>>> attention_mask[:, -1]
tensor([0, 1], device='cuda:0')

If there was indeed no padding we would simply select the last token embedding for each of the examples with last_hidden_states[:, -1]. However, since we have padding we need to select the last non-padding token embedding from each example in the batch. In order to pick this embedding we need to get its index for each example. This is achieved via

>>> sequence_lengths = attention_mask.sum(dim=1) - 1
>>> sequence_lengths
tensor([3, 6], device='cuda:0')

So now we need to simply index into the tensor, with the correct indices in the first two dimensions. To get the indices for all the examples in the batch we can use torch.arange as follows:

>>> torch.arange(batch_size, device=last_hidden_states.device)
tensor([0, 1], device='cuda:0')

Then we can pluck out the correct token embeddings for each example using this and the indices of the last non padding token:

>>> embeddings = last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
>>> embeddings.shape
torch.Size([2, 3584])

And we get two embeddings for the two examples passed in!

How to Run — TLDR

The full code separated out into functions looks like

import numpy as np
import numpy.typing as npt
import torch
import transformersDEVICE = torch.device("cuda")
def last_token_pool(last_hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
# checks whether there is any padding (where attention mask = 0 for a given text)
no_padding = attention_mask[:, -1].sum() == attention_mask.shape[0]
# if no padding - only would happen if batch size of 1 or all sequnces have the same length, then take the last tokens as the embeddings
if no_padding:
return last_hidden_states[:, -1]
# otherwise use the last non padding token for each text in the batch
sequence_lengths = attention_mask.sum(dim=1) - 1
batch_size = last_hidden_states.shape[0]
return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths
def encode_with_qwen_model(
model: transformers.PreTrainedModel,
tokenizer: transformers.tokenization_utils.PreTrainedTokenizer | transformers.tokenization_utils_fast.PreTrainedTokenizerFast,
texts: list[str],
max_length: int = 32768,
) -> npt.NDArray[np.float16]:
batch_dict = tokenizer(texts, max_length=max_length, padding=True, truncation=True, return_tensors="pt").to(DEVICE)
with torch.no_grad():
outputs = model(**batch_dict)
embeddings = last_token_pool(outputs.last_hidden_state, batch_dict["attention_mask"])
return embeddings.cpu().numpy()
def main() -> None:
model_path = "Alibaba-NLP/gte-Qwen2-7B-instruct"
tokenizer = transformers.AutoTokenizer.from_pretrained(model_path)
model = transformers.AutoModel.from_pretrained(model_path, trust_remote_code=True, torch_dtype=torch.float16).to(DEVICE)
print("Loaded tokeniser and model")
texts_to_encode = ["example text 1", "example text 2 of different length"]
embeddings = encode_with_qwen_model(model, tokenizer, texts_to_encode)
print(embeddings.shape)
if __name__ == "__main__":
main()

The encode_with_qwen_model returns a numpy array. In order to convert a PyTorch tensor to a numpy array we first have to move it off the GPU back onto the CPU which is achieved with the cpu() method. Please note that if you are planning to run long texts you should reduce the batch size to 1 and only embed one example at a time (thus reducing the list texts_to_encode to length 1).