Ejecución de un modelo de incorporación de parámetros SOTA 7B en una sola GPU | de Szymon Palucha | agosto de 2024

Configuración

El modelo con el que experimentaremos es el Alibaba-NLP/gte-Qwen2-7B-instruct De Transformers. La tarjeta modelo es aquí.

Para realizar este experimento, he utilizado Python 3.10.8 e instalado los siguientes paquetes:

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

Me encontré con algunos dificultad En la instalación flash-attn Se requiere ejecutar este modelo y, por lo tanto, tuve que instalar la versión específica que se detalla arriba. Si alguien tiene una mejor solución, ¡háganmelo saber!

La instancia de Amazon SageMaker que utilicé para este experimento es la ml.g5.2xlargeTiene una GPU NVIDIA A10G de 24 GB y 32 GB de memoria de CPU y cuesta 1,69 USD por hora. La siguiente captura de pantalla de AWS muestra todos los detalles de la instancia.

En realidad, para ser precisos, si corres nvidia-smi Verás que la instancia solo tiene 23 GB de memoria de GPU, que es un poco menos de lo anunciado. La versión de CUDA en esta GPU es 12.2.

Cómo correr: en detalle

Si observa la tarjeta modelo, una de las formas sugeridas de usar este modelo es a través de sentence-transformers biblioteca como se muestra a continuación

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 es una extensión del paquete Transformers para calcular incrustaciones y es muy útil ya que puede hacer que las cosas funcionen con dos líneas de código. La desventaja es que tiene menos control sobre cómo cargar el modelo, ya que oculta los detalles de tokenización y agrupación. El código anterior No correrá en nuestra instancia de GPU porque intenta cargar el modelo con precisión float32 completa, lo que ocuparía 28 GB de memoria. Cuando se inicializa el modelo del transformador de oraciones, este verifica si hay dispositivos disponibles (cuda para GPU) y automáticamente cambia el modelo de Pytorch al dispositivo. Como resultado, se bloquea después de cargar 5/7 del modelo y se bloquea.

En lugar de eso, necesitamos poder cargar el modelo con precisión float16 antes de moverlo a la GPU. Por lo tanto, necesitamos usar la biblioteca Transformers de nivel inferior. (No estoy seguro de cómo hacerlo con los transformadores de oraciones, ¡pero avísenme si existe alguno!) Lo hacemos de la siguiente manera

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

Con el torch_dtype parámetro especificamos que el modelo debe cargarse con precisión float16 de inmediato, por lo que solo se requieren 14 GB de memoria. Luego, debemos mover el modelo al dispositivo GPU, lo que se logra con el to Método. ¡Usando el código anterior, el modelo tarda casi 2 minutos en cargarse!

Ya que estamos usando transformers Necesitamos cargar por separado el tokenizador para tokenizar los textos de entrada de la siguiente manera:

tokenizer = transformers.AutoTokenizer.from_pretrained(model_path)

El siguiente paso es tokenizar los textos de entrada, lo que se hace de la siguiente manera:

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)

La longitud máxima del modelo Qwen2 es 32678, sin embargo, como veremos más adelante, no podemos ejecutarlo con una secuencia tan larga en nuestra GPU de 24 GB debido a los requisitos de memoria adicionales. Recomendaría reducir esto a no más de 24.000 para evitar errores de falta de memoria. El relleno garantiza que todas las entradas del lote tengan la misma longitud, mientras que el truncamiento garantiza que se truncarán todas las entradas que superen la longitud máxima. Para obtener más información, consulte la documentosPor último, nos aseguramos de devolver los tensores de PyTorch (el valor predeterminado serían listas) y mover estos tensores a la GPU para que estén disponibles para pasarlos al modelo.

El siguiente paso es pasar las entradas a través de nuestro modelo y realizar la agrupación. Esto se hace de la siguiente manera

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

con el last_token_pool que se ve así:

def last_token_pool(last_hidden_states: antorcha.Tensor, máscara_de_atención: antorcha.Tensor) -> antorcha.Tensor:
# comprueba si hay relleno (donde la máscara de atención = 0 para un texto determinado)
sin relleno = máscara de atención[:, -1].sum() == máscara_de_atención.forma[0]
# si no hay relleno - solo sucedería si el tamaño del lote es 1 o todas las secuencias tienen la misma longitud, luego se toman los últimos tokens como incrustaciones
si no hay relleno:
devolver los últimos estados ocultos[:, -1]
# de lo contrario, use el último token sin relleno para cada texto en el lote
longitudes_de_secuencia = máscara_de_atención.suma(dim=1) - 1
tamaño_lote = últimos_estados_ocultos.forma[0]
devolver los últimos estados ocultos[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] Sin embargo, el modelo Qwen2 se basa en LLM, es decir, en un decodificador de transformadores. En el decodificador, los tokens se generan de forma autorregresiva (uno tras otro) y, por lo tanto, el último token contiene toda la información codificada sobre la oración.
• El objetivo de la last_token_pool La función es, por lo tanto, seleccionar la incrustación del último token generado (que no era el token de relleno) para cada ejemplo del lote.
• Utiliza el attention_mask que le dice al modelo cuáles de los tokens son tokens de relleno para cada ejemplo en el lote (ver la documentos).

Ejemplo anotado

Veamos un ejemplo para entenderlo con un poco más de detalle. Supongamos que queremos incrustar dos ejemplos en un solo lote:

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

Las salidas del tokenizador (el batch_dict ) se verá de la siguiente manera:

>>> 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')}

De esto se puede ver que la primera oración se divide en cuatro tokens (8687, 1467, 220, 16), mientras que la segunda oración se divide en siete tokens. Como resultado, la primera oración se rellena (con tres tokens de relleno con id 151643) hasta la longitud siete, el máximo en el lote. La máscara de atención refleja esto: tiene tres ceros para el primer ejemplo que corresponden a la ubicación de los tokens de relleno. Ambos tensores tienen el mismo tamaño.

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

Ahora pasando el batch_dict A través del modelo podemos recuperar el último estado oculto de la forma del modelo:

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

Podemos ver que esto tiene forma. (tamaño de lote, longitud máxima de secuencia, dimensión de incrustación)¡Qwen2 tiene una dimensión de incrustación de 3584!

Ahora estamos en el last_token_pool función. La primera línea verifica si existe relleno, lo hace sumando la última “columna” de attention_mask y comparándola con el tamaño del lote (dado por attention_mask.shape[0]Esto solo resultará verdadero si existe un 1 en todas las máscaras de atención, es decir, si todos los ejemplos tienen la misma longitud o si solo tenemos un ejemplo.

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

Si de hecho no hubiera relleno, simplemente seleccionaríamos la última incrustación de token para cada uno de los ejemplos con last_hidden_states[:, -1]Sin embargo, dado que tenemos relleno, necesitamos seleccionar la última incrustación de token sin relleno de cada ejemplo del lote. Para elegir esta incrustación, necesitamos obtener su índice para cada ejemplo. Esto se logra mediante

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

Ahora simplemente necesitamos indexar el tensor con los índices correctos en las dos primeras dimensiones. Para obtener los índices de todos los ejemplos del lote, podemos usar torch.arange como sigue:

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

Luego podemos extraer las incrustaciones de token correctas para cada ejemplo usando esto y los índices del último token sin relleno:

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

¡Y obtenemos dos incrustaciones para los dos ejemplos pasados!

Cómo correr: resumen

El código completo separado en funciones se ve así

import numpy as np
import numpy.typing as npt
import torch
import transformers
DEVICE = 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()

El encode_with_qwen_model devuelve una matriz numpy. Para convertir un tensor de PyTorch en una matriz numpy, primero tenemos que sacarlo de la GPU y volverlo a colocar en la CPU, lo que se logra con el comando cpu() método. Tenga en cuenta que si planea ejecutar textos largos, debe reducir el tamaño del lote a 1 y solo incrustar un ejemplo a la vez (reduciendo así la lista). texts_to_encode hasta longitud 1).