Desplegar Falcon-40B con DLCs de inferencia de modelos grandes en Amazon SageMaker.

Deploy Falcon-40B with large model inference DLCs on Amazon SageMaker.

La semana pasada, Technology Innovation Institute (TII) lanzó TII Falcon LLM, un modelo lingüístico de gran tamaño (LLM) fundamental de código abierto. Entrenado con 1 billón de tokens con Amazon SageMaker, Falcon cuenta con un rendimiento de primera categoría (número 1 en el ranking de Hugging Face al momento de escribir este artículo), al mismo tiempo que es comparativamente ligero y menos costoso de alojar que otros LLM como llama-65B. En esta publicación, demostramos cómo desplegar Falcon para aplicaciones como la comprensión del lenguaje y la asistencia de escritura automatizada utilizando contenedores de aprendizaje profundo de inferencia de modelos grandes en SageMaker.

El Falcon ha aterrizado en SageMaker

TII es la organización de investigación aplicada dentro del Consejo de Investigación de Tecnología Avanzada de Abu Dhabi; su equipo de científicos, investigadores e ingenieros está dedicado al descubrimiento de tecnologías transformadoras y al desarrollo de avances científicos que protejan nuestro futuro. A principios de este año, TII se propuso entrenar un LLM de código abierto de última generación y utilizó la infraestructura, las herramientas y la experiencia de SageMaker para hacer el trabajo (para obtener más información sobre cómo se entrenó este modelo en SageMaker, consulte la publicación Technology Innovation Institute entrena el modelo fundamental Falcon LLM 40B en Amazon SageMaker). El resultado de este esfuerzo es TII Falcon LLM.

Entrenado con 1 billón de tokens, Falcon cuenta con un rendimiento de primera categoría frente a la Evaluación de Arneses de Lenguaje de IA de Eleuther y ocupa actualmente el primer lugar en el ranking de Hugging Face por precisión. El modelo está disponible en dos tamaños diferentes: Falcon-40B y Falcon-7B, y se puede utilizar para obtener un rendimiento de última generación en aplicaciones como la comprensión del lenguaje, las experiencias conversacionales y la asistencia de escritura automatizada. Esta publicación le ayudará a comenzar a desplegar Falcon en SageMaker para obtener una inferencia de alta precisión en estos tipos de dominios.

Los contenedores de aprendizaje profundo de inferencia de modelos grandes de SageMaker simplifican el alojamiento de LLM

El alojamiento de LLM como Falcon-40B y Falcon-7B puede ser un desafío. Los modelos más grandes suelen ser más precisos porque incluyen miles de millones de parámetros, pero su tamaño también puede provocar una latencia de inferencia más lenta o un peor rendimiento. El alojamiento de un LLM puede requerir más memoria de GPU y núcleos optimizados para lograr un rendimiento aceptable. Para complicar aún más las cosas, aunque los modelos más pequeños como Falcon-7B generalmente pueden caber en una sola GPU como una instancia NVIDIA A10G que alimenta los tipos de instancia AWS G5, los modelos más grandes como Falcon-40B no pueden hacerlo. Cuando esto sucede, se deben utilizar estrategias como el paralelismo de tensor para dividir ese modelo más grande en múltiples piezas y aprovechar la memoria de múltiples GPU. Las soluciones de alojamiento heredadas utilizadas para modelos más pequeños generalmente no ofrecen este tipo de funcionalidad, lo que aumenta la dificultad.

Los contenedores de aprendizaje profundo de inferencia de modelos grandes (LMI DLC) de SageMaker pueden ayudar. LMI DLC es una solución completa de extremo a extremo para alojar LLM como Falcon-40B. En el extremo frontal, incluyen un servidor de modelo de alto rendimiento (DJL Serving) diseñado para inferencia de modelos grandes con características como transmisión de tokens y replicación automática del modelo dentro de una instancia para aumentar el rendimiento. En la parte posterior, LMI DLC también incluye varios motores de paralelismo de modelo de alto rendimiento, como DeepSpeed y FasterTransformer, que pueden dividir y administrar los parámetros del modelo en múltiples GPU. Estos motores también incluyen núcleos optimizados para modelos de transformador populares, lo que puede acelerar la inferencia hasta tres veces más rápido. Con LMI DLC, solo es necesario crear un archivo de configuración para comenzar a alojar LLM en SageMaker. Para obtener más información sobre LMI DLC de SageMaker, consulte el paralelismo de modelo y la inferencia de modelos grandes y nuestra lista de imágenes disponibles. También puede ver nuestra publicación anterior sobre el alojamiento de Bloom-175B en SageMaker utilizando LMI DLC.

Descripción general de la solución

Esta publicación le guía sobre cómo alojar Falcon-40B utilizando DeepSpeed en SageMaker mediante LMI DLC. Falcon-40B requiere que usemos múltiples GPU A10, mientras que Falcon-7B solo requiere una GPU. También hemos preparado ejemplos a los que puede hacer referencia para alojar Falcon-40B y Falcon-7B utilizando tanto DeepSpeed como Accelerate. Puede encontrar nuestros ejemplos de código en GitHub.

Este ejemplo se puede ejecutar en instancias de notebook de SageMaker o en cuadernos de Amazon SageMaker Studio. Para alojar Falcon-40B utilizando LMI y DeepSpeed, debemos usar una instancia ml.g5.24xlarge. Estas instancias proporcionan 4 GPU NVIDIA A10G, cada una de las cuales admite 96 GiB de memoria de GPU. Además, el host proporciona 96 vCPU y 384 GiB de memoria del host. El contenedor LMI ayudará a abordar gran parte del trabajo pesado no diferenciado asociado con el alojamiento de LLM, incluida la descarga del modelo y la partición del artefacto del modelo para que sus parámetros constituyentes se puedan distribuir en varias GPU.

Las cuotas para las instancias de aprendizaje automático (ML) de SageMaker pueden variar entre cuentas. Si recibe un error que indica que ha superado su cuota de instancias g5.24xlarge mientras sigue esta publicación, puede aumentar el límite a través de la consola de Cuotas de Servicio.

Paseo por la libreta

Para comenzar, comenzamos instalando e importando las dependencias necesarias para nuestro ejemplo. Usamos el SDK de Boto3 y el SDK de SageMaker. Tenga en cuenta que usamos Amazon Simple Storage Service (Amazon S3) para almacenar los artefactos del modelo que necesitamos para que SageMaker y LMI los usen, por lo que configuramos una variable de prefijo S3 en consecuencia. Vea el siguiente código:

import sagemaker
import jinja2
from sagemaker import image_uris
import boto3
import os
import time
import json
from pathlib import Path
from sagemaker.utils import name_from_base

role = sagemaker.get_execution_role()  # papel de ejecución para el endpoint
sess = sagemaker.session.Session()  # sesión de SageMaker para interactuar con diferentes APIs de AWS
bucket = sess.default_bucket()  # cubo para alojar artefactos
model_bucket = sess.default_bucket()  # cubo para alojar artefactos
s3_code_prefix_deepspeed = "hf-large-model-djl-/code_falcon40b/deepspeed"  # carpeta dentro del cubo donde irá el artefacto del código
region = sess._region_name
account_id = sess.account_id()
s3_client = boto3.client("s3")
sm_client = boto3.client("sagemaker")
smr_client = boto3.client("sagemaker-runtime")
jinja_env = jinja2.Environment()

Luego creamos una carpeta local para nuestro espacio de trabajo para almacenar los artefactos de nuestro modelo:

!mkdir -p code_falcon40b_deepspeed

Primero creamos un archivo de configuración serving.properties en el directorio local que creamos. Este archivo serving.properties indica al contenedor LMI y a la biblioteca DJL Serving del front-end qué motor de paralelización de modelos y optimización de inferencia queremos usar. Puede encontrar las opciones de configuración tanto para DeepSpeed como para Hugging Face Accelerate en Configuraciones y ajustes. Aquí, tenga en cuenta que establecemos el parámetro option.model_id para definir qué modelo de Hugging Face queremos extraer. SageMaker simplifica el trabajo con modelos de Hugging Face, y esta línea es todo lo que necesita. Además, establecemos option.tensor_parallel_degree en un valor de 4 porque tenemos cuatro GPUs en nuestra instancia ml.g5.24xlarge. Este parámetro define cuántas particiones del modelo crear y distribuir. Tenga en cuenta que si hubiéramos utilizado una instancia más grande con ocho GPUs, como ml.g5.48xlarge, y aún así estableciéramos un valor de 4, entonces LMI crearía automáticamente dos réplicas del modelo (dos réplicas distribuidas en cuatro GPUs cada una). Vea el siguiente código:

%%writefile ./code_falcon40b_deepspeed/serving.properties
engine=Python
#para implementar falcon-40b-instruct, establezca el valor de model_id en 'tiiuae/falcon-40b-instruct'
option.model_id=tiiuae/falcon-40b
option.tensor_parallel_degree=4
#option.s3url = {{s3url}}

También puede cambiar tiiuae/falcon-40b por tiiuae/falcon-40b-instruct si se adapta mejor a sus necesidades.

También incluimos un archivo requirements.txt que puede especificar para instalar paquetes que requiera:

%%writefile ./code_falcon40b_deepspeed/requirements.txt
einops
torch==2.0.1

Lo último que necesitamos es el archivo model.py que se utilizará con su modelo:

%%writefile ./code_falcon40b_deepspeed/model.py
from djl_python import Input, Output
import os
import torch
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
from typing import Any, Dict, Tuple
import warnings

predictor = None


def get_model(properties):
    model_name = properties["model_id"]
    local_rank = int(os.getenv("LOCAL_RANK", "0"))
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        low_cpu_mem_usage=True,
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
        device_map="auto",
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    generator = pipeline(
        task="text-generation", model=model, tokenizer=tokenizer, device_map="auto"
    )
    return generator


def handle(inputs: Input) -> None:
    global predictor
    if not predictor:
        predictor = get_model(inputs.get_properties())
    if inputs.is_empty():
        # El servidor de modelos realiza una llamada vacía para calentar el modelo al iniciarse
        return None
    data = inputs.get_as_json()
    text = data["text"]
    text_length = data["text_length"]
    outputs = predictor(text, do_sample=True, min_length=text_length, max_length=text_length)
    result = {"outputs": outputs}
    return Output().add_as_json(result)

¡Eso es! En este punto, hemos creado todos los artefactos que necesitarás para implementar Falcon-40B con DeepSpeed. Empaquetamos el directorio en un archivo *.tar.gz y lo subimos a Amazon S3. Ten en cuenta que el modelo real no se ha descargado ni empaquetado en este archivo. El contenedor LMI descargará el modelo para ti directamente desde Hugging Face. También tienes la opción de apuntar a un bucket de S3 si deseas tener tu propia copia del modelo en una ubicación más adecuada para descargar. LMI también incluye optimización para descargas desde Amazon S3 con alta velocidad de rendimiento. Consulta el siguiente código:

s3_code_artifact_deepspeed= sess.upload_data("model.tar.gz", bucket, s3_code_prefix_deepspeed)
print(f"El código S3 o el archivo tar del modelo para deepspeed se cargó en --- > {s3_code_artifact_deepspeed}")

Todo lo que queda por hacer en este punto es definir el contenedor que queremos usar y crear un objeto de modelo:

inference_image_uri = (
    f"763104351884.dkr.ecr.{region}.amazonaws.com/djl-inference:0.22.1-deepspeed0.8.3-cu118"
)
model_name_acc = name_from_base(f"falcon40b-model-ds")
create_model_response = sm_client.create_model(
    ModelName=model_name_acc,
    ExecutionRoleArn=role,
    PrimaryContainer={"Image": inference_image_uri, "ModelDataUrl": s3_code_artifact_deepspeed},
)
model_arn = create_model_response["ModelArn"]

Luego creamos una configuración de punto de enlace y creamos el punto de enlace:

endpoint_config_name = f"{model_name}-config"
endpoint_name = f"{model_name}-endpoint"
endpoint_config_response = sm_client.create_endpoint_config(
    EndpointConfigName=endpoint_config_name,
    ProductionVariants=[
        {
            "VariantName": "variant1",
            "ModelName": model_name,
            "InstanceType": "ml.g5.24xlarge",
            "InitialInstanceCount": 1,
            "ModelDataDownloadTimeoutInSeconds": 3600,
            "ContainerStartupHealthCheckTimeoutInSeconds": 3600,
            # "VolumeSizeInGB": 512
        },
    ],
)
endpoint_config_response

create_endpoint_response = sm_client.create_endpoint(
    EndpointName=f"{endpoint_name}", EndpointConfigName=endpoint_config_name
)
print(f"Se ha creado el punto de enlace: {create_endpoint_response['EndpointArn']}")

Elementos de configuración a tener en cuenta para una implementación exitosa

Una consideración importante para la implementación de modelos grandes es asegurar que haya tiempo suficiente para descargar el modelo de Hugging Face. En nuestras pruebas, Falcon-40B tardó alrededor de 90 minutos en descargarse en la instancia. Un conjunto clave de configuraciones para permitir esto son ContainerStartupHealthCheckTimeoutInSeconds y ModelDataDownloadTimeoutInSeconds. Asegúrate de que la configuración de punto de enlace de SageMaker tenga un valor de 3600 para cada uno de estos. Además, es mucho más fácil descargar desde Amazon S3 en lugar del modelo original de zoológico utilizando los contenedores LMI que están especialmente diseñados para LLMS que utilizan la utilidad S5cmd, lo que reduce el tiempo de descarga del modelo a alrededor de 10 minutos.

Puedes monitorear el estado del punto de enlace llamando a DescribeEndpoint, lo que te indicará cuándo todo esté completo. ¡Tu punto de enlace ahora está listo para responder a solicitudes de inferencia! Debido a que LMI maneja la partición y orquestación del modelo por ti, cada solicitud se procesará utilizando las 4 GPU disponibles en nuestra instancia ml.g5.12xlarge. Esto nos permite alojar LLM y aumentar el rendimiento si se escalan los aceleradores de GPU horizontalmente. Consulta el siguiente código:

response_model = smr_client.invoke_endpoint(
    EndpointName=endpoint_name,
    Body=json.dumps({"text": "¿Cuál es el propósito de la vida?", "text_length": 150}),
    ContentType="application/json",
)

response_model["Body"].read().decode("utf8")

Si has terminado y deseas eliminar la configuración del punto de enlace, el punto de enlace y el objeto de modelo, puedes ejecutar los siguientes comandos:

sm_client.delete_endpoint(EndpointName=endpoint_name)
sm_client.delete_endpoint_config(EndpointConfigName=endpoint_config_name)
sm_client.delete_model(ModelName=model_name)

Este código al que hacemos referencia en esta publicación se puede encontrar en el cuaderno completo en GitHub.

Conclusión

SageMaker Hosting y el LMI DLC hacen que sea fácil para ti alojar LLM como Falcon-40B. Se encarga del trabajo pesado no diferenciado en la orquestación de lo que se requiere para alojar modelos en varias GPU y proporciona opciones configurables para adaptarse a tus necesidades. Además, el uso de modelos de Hugging Face se vuelve muy sencillo, con soporte integrado para estos modelos.

En este post, mostramos cómo puedes utilizar SageMaker para alojar el modelo Falcon-40B utilizando DeepSpeed. Además, proporcionamos ejemplos en GitHub para alojar Falcon-40B utilizando Accelerate, y los modelos más pequeños Falcon-7B. ¡Te animamos a que lo pruebes en SageMaker con LMI y te pongas manos a la obra con el mejor LLM disponible públicamente hasta la fecha!