Implementar GPT-J 6B para inferencia utilizando Hugging Face Transformers y Amazon SageMaker

Implementar GPT-J 6B para inferencia con Hugging Face Transformers y Amazon SageMaker.

Hace casi 6 meses, EleutherAI lanzó GPT-J 6B, una alternativa de código abierto a GPT-3 de OpenAI. GPT-J 6B es el sucesor de 6 mil millones de parámetros de la familia GPT-NEO de EleutherAIs, una familia de modelos de lenguaje basados en transformadores y basados en la arquitectura GPT para la generación de texto.

El objetivo principal de EleutherAI es entrenar un modelo del mismo tamaño que GPT-3 y ponerlo a disposición del público bajo una licencia abierta.

En los últimos 6 meses, GPT-J ha despertado mucho interés entre investigadores, científicos de datos e incluso desarrolladores de software, pero ha sido muy desafiante implementar GPT-J en producción para casos de uso y productos del mundo real.

Existen algunas soluciones alojadas para utilizar GPT-J en cargas de trabajo de producción, como la API de inferencia de Hugging Face, o para experimentar utilizando el playground 6b de EleutherAI, pero hay menos ejemplos de cómo implementarlo fácilmente en tu propio entorno.

En esta publicación de blog, aprenderás cómo implementar fácilmente GPT-J utilizando Amazon SageMaker y el kit de herramientas de inferencia de Hugging Face con unas pocas líneas de código para obtener inferencia en tiempo real escalable, confiable y segura utilizando una instancia de GPU de tamaño regular con NVIDIA T4 (~500$/mes).

Pero antes de entrar en eso, quiero explicar por qué es desafiante implementar GPT-J en producción.


Antecedentes

Los pesos del modelo de 6 mil millones de parámetros representan una huella de memoria de ~24GB. Para cargarlo en float32, se necesitaría al menos el doble de la memoria RAM de tamaño del modelo: 1x para los pesos iniciales y otro 1x para cargar el punto de control. Por lo tanto, para GPT-J se requerirían al menos 48GB de memoria RAM de la CPU solo para cargar el modelo.

Para hacer el modelo más accesible, EleutherAI también proporciona pesos en float16, y transformers tiene nuevas opciones para reducir la huella de memoria al cargar modelos de lenguaje grandes. Combinando todo esto, debería tomar aproximadamente 12.1GB de memoria RAM de la CPU para cargar el modelo.

from transformers import GPTJForCausalLM
import torch

model = GPTJForCausalLM.from_pretrained(
    "EleutherAI/gpt-j-6B",
        revision="float16",
        torch_dtype=torch.float16,
        low_cpu_mem_usage=True
)

La advertencia de este ejemplo es que lleva mucho tiempo cargar el modelo en memoria y estar listo para su uso. En mis experimentos, tardó 3 minutos y 32 segundos en cargar el modelo con el fragmento de código anterior en una instancia de AWS EC2 P3.2xlarge (el modelo no se almacenaba en disco). Esta duración se puede reducir almacenando el modelo ya en disco, lo que reduce el tiempo de carga a 1 minuto y 23 segundos, lo cual sigue siendo muy largo para cargas de trabajo de producción donde se necesita considerar la escalabilidad y la confiabilidad.

Por ejemplo, Amazon SageMaker tiene un límite de 60 segundos para responder a las solicitudes, lo que significa que el modelo debe cargarse y las predicciones deben ejecutarse dentro de los 60 segundos, lo cual tiene mucho sentido para mantener el modelo/endpoint escalable y confiable para tu carga de trabajo. Si tienes predicciones más largas, puedes usar batch-transform.

En Transformers, los modelos cargados con el método from_pretrained siguen la práctica recomendada de PyTorch, lo que lleva alrededor de 1.97 segundos para BERT [REF]. PyTorch ofrece una forma adicional de guardar y cargar modelos utilizando torch.save(model, PATH) y torch.load(PATH).

“Guardar un modelo de esta manera guardará todo el módulo utilizando el módulo pickle de Python. La desventaja de este enfoque es que los datos serializados están vinculados a las clases específicas y la estructura de directorios exacta utilizada cuando se guarda el modelo.”

Esto significa que cuando guardamos un modelo con transformers==4.13.2, podría ser potencialmente incompatible al intentar cargarlo con transformers==4.15.0. Sin embargo, cargar modelos de esta manera reduce el tiempo de carga en ~12x, hasta 0.166s para BERT.

Aplicando esto a GPT-J significa que podemos reducir el tiempo de carga de 1 minuto y 23 segundos a 7.7 segundos, lo cual es ~10.5x más rápido.

Figura 1. Tiempo de carga del modelo de BERT y GPTJ

Tutorial

Con este método de guardar y cargar modelos, logramos un rendimiento de carga del modelo de GPT-J compatible con escenarios de producción. Pero debemos tener en cuenta que debemos alinear:

Alinear la versión de PyTorch y Transformers al guardar el modelo con torch.save(model, PATH) y cargar el modelo con torch.load(PATH) para evitar incompatibilidades.

Guardar GPT-J usando torch.save

Para crear nuestro archivo de modelo compatible con torch.load(), cargamos GPT-J usando Transformers y el método from_pretrained, y luego lo guardamos con torch.save().

from transformers import AutoTokenizer,GPTJForCausalLM
import torch

# cargar el modelo fp 16
model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B", revision="float16", torch_dtype=torch.float16)
# guardar el modelo con torch.save
torch.save(model, "gptj.pt")

Ahora podemos cargar nuestro modelo de GPT-J con torch.load() para realizar predicciones.

from transformers import pipeline
import torch

# cargar el modelo
model = torch.load("gptj.pt")
# cargar el tokenizador
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")

# crear el pipeline
gen = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)

# realizar la predicción
gen("Mi nombre es philipp")
#[{'generated_text': 'Mi nombre es philipp k. y vivo justo fuera de Detroit....

Crear model.tar.gz para el endpoint en tiempo real de Amazon SageMaker

Dado que podemos cargar nuestro modelo rápidamente y realizar inferencias en él, vamos a implementarlo en Amazon SageMaker.

Hay dos formas de implementar transformers en Amazon SageMaker. Puedes “Implementar un modelo desde el Hugging Face Hub” directamente o “Implementar un modelo con model_data almacenado en S3″. Dado que no estamos utilizando el método predeterminado de Transformers, debemos optar por la segunda opción y implementar nuestro endpoint con el modelo almacenado en S3.

Para esto, necesitamos crear un artefacto model.tar.gz que contenga los pesos de nuestro modelo y los archivos adicionales que necesitamos para la inferencia, como tokenizer.json.

Proporcionamos artefactos model.tar.gz subidos y accesibles públicamente, que se pueden utilizar con HuggingFaceModel para implementar GPT-J en Amazon SageMaker.

Consulta “Implementar GPT-J como endpoint de Amazon SageMaker” para saber cómo utilizarlos.

Si aún quieres o necesitas crear tu propio model.tar.gz, por ejemplo, debido a pautas de cumplimiento, puedes utilizar el script de ayuda convert_gpt.py para este propósito, que crea el model.tar.gz y lo carga en S3.

# clonar el directorio
git clone https://github.com/philschmid/amazon-sagemaker-gpt-j-sample.git

# cambiar al directorio amazon-sagemaker-gpt-j-sample
cd amazon-sagemaker-gpt-j-sample

# crear y cargar model.tar.gz
pip3 install -r requirements.txt
python3 convert_gptj.py --bucket_name {model_storage}

El convert_gpt.py debería imprimir una URI de S3 similar a esta. s3://hf-sagemaker-inference/gpt-j/model.tar.gz.

Implementar GPT-J como endpoint de Amazon SageMaker

Para implementar nuestro endpoint de Amazon SageMaker, vamos a utilizar el SDK de Python de Amazon SageMaker y la clase HuggingFaceModel.

El fragmento de código a continuación utiliza get_execution_role, que solo está disponible dentro de las instancias de bloc de notas o Studio de Amazon SageMaker. Si quieres implementar un modelo fuera de él, consulta la documentación.

El model_uri define la ubicación de nuestro artefacto del modelo GPT-J. Vamos a utilizar el que está disponible públicamente proporcionado por nosotros.

from sagemaker.huggingface import HuggingFaceModel
import sagemaker

# IAM role con permisos para crear endpoint
role = sagemaker.get_execution_role()

# URI público de S3 al artefacto de gpt-j
model_uri="s3://huggingface-sagemaker-models/transformers/4.12.3/pytorch/1.9.1/gpt-j/model.tar.gz"

# crear la clase de modelo de Hugging Face
huggingface_model = HuggingFaceModel(
    model_data=model_uri,
    transformers_version='4.12.3',
    pytorch_version='1.9.1',
    py_version='py38',
    role=role, 
)

# desplegar el modelo en SageMaker Inference
predictor = huggingface_model.deploy(
    initial_instance_count=1, # número de instancias
    instance_type='ml.g4dn.xlarge' #'ml.p3.2xlarge' # tipo de instancia ec2
)

Si desea utilizar su propio model.tar.gz, simplemente reemplace el model_uri con su URI de S3.

La implementación debería tomar alrededor de 3-5 minutos.

Ejecutar predicciones

Podemos ejecutar predicciones utilizando las instancias predictor creadas por nuestro método .deploy. Para enviar una solicitud a nuestro endpoint, usamos predictor.predict con nuestros inputs.

predictor.predict({
    "inputs": "¿Podría proporcionarnos más detalles sobre su "
})

Si desea personalizar sus predicciones utilizando argumentos adicionales como min_length, consulte las “Mejores prácticas de uso” a continuación.

Mejores prácticas de uso

Cuando se utilizan modelos generativos, la mayoría de las veces se desea configurar o personalizar la predicción para adaptarla a sus necesidades, por ejemplo, utilizando la búsqueda de haz, configurando la longitud máxima o mínima de la secuencia generada o ajustando la temperatura para reducir la repetición. La biblioteca Transformers proporciona diferentes estrategias y kwargs para hacer esto, el toolkit de Inferencia de Hugging Face ofrece la misma funcionalidad utilizando el atributo parameters de su carga útil de solicitud. A continuación, puede encontrar ejemplos de cómo generar texto sin parámetros, con búsqueda de haz y utilizando configuraciones personalizadas. Si desea aprender sobre diferentes estrategias de decodificación, consulte esta publicación de blog.

Solicitud predeterminada

Este es un ejemplo de una solicitud predeterminada utilizando la búsqueda greedy.

Tiempo de inferencia después de la primera solicitud: 3s

predictor.predict({
    "inputs": "¿Podría proporcionarnos más detalles sobre su "
})

Solicitud de búsqueda de haz

Este es un ejemplo de una solicitud utilizando la búsqueda de haz con 5 haces.

Tiempo de inferencia después de la primera solicitud: 3.3s

predictor.predict({
    "inputs": "¿Podría proporcionarnos más detalles sobre su ",
  "parameters" : {
    "num_beams": 5,
  }
})

Solicitud con parámetros

Este es un ejemplo de una solicitud utilizando un parámetro personalizado, por ejemplo, min_length para generar al menos 512 tokens.

Tiempo de inferencia después de la primera solicitud: 38s

predictor.predict({
    "inputs": "¿Podría proporcionarnos más detalles sobre su ",
  "parameters" : {
    "max_length": 512,
    "temperature": 0.9,
  }
})

Ejemplo de pocos datos (avanzado)

Este es un ejemplo de cómo podría utilizar eos_token_id para detener la generación en un cierto token, por ejemplo, \n, . o ### para predicciones de pocos datos. A continuación se muestra un ejemplo de pocos datos para generar tweets para palabras clave.

Tiempo de inferencia después de la primera solicitud: 15-45s

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")

end_sequence="###"
temperature=4
max_generated_token_length=25
prompt= """key: mercados
tweet: Toma retroalimentación de la naturaleza y los mercados, no de las personas.
###
key: niños
tweet: Tal vez morimos para poder volver como niños.
###
key: startups
tweet: Las startups no deberían preocuparse por cómo apagar incendios, deberían preocuparse por cómo encenderlos.
###
key: hugging face
tweet:"""

predictor.predict({
    'inputs': prompt,
  "parameters" : {
    "max_length": int(len(prompt) + max_generated_token_length),
    "temperature": float(temperature),
    "eos_token_id": int(tokenizer.convert_tokens_to_ids(end_sequence)),
    "return_full_text":False
  }
})

Para eliminar tu endpoint, puedes ejecutar:

predictor.delete_endpoint()

Conclusión

Logramos implementar correctamente GPT-J, un modelo de lenguaje de 6 mil millones de parámetros creado por EleutherAI, utilizando Amazon SageMaker. Redujimos el tiempo de carga del modelo de 3.5 minutos a 8 segundos para poder ejecutar inferencias escalables y confiables.

Recuerda que el uso de torch.save() y torch.load() puede generar problemas de compatibilidad. Si deseas obtener más información sobre cómo escalar tus endpoints de Amazon SageMaker, consulta mi otro artículo de blog: “MLOps: End-to-End Hugging Face Transformers con el Hub y las canalizaciones de SageMaker”.


¡Gracias por leer! Si tienes alguna pregunta, no dudes en contactarme a través de Github o en el foro. También puedes conectarte conmigo en Twitter o LinkedIn.