Optimiza el costo de implementación de los modelos base de Amazon SageMaker JumpStart con los puntos finales asincrónicos de Amazon SageMaker

Optimiza costos de implementación de modelos base de Amazon SageMaker JumpStart con puntos finales asincrónicos de Amazon SageMaker

El éxito de las aplicaciones de IA generativa en una amplia gama de industrias ha atraído la atención e interés de empresas de todo el mundo que buscan reproducir y superar los logros de sus competidores o resolver casos de uso nuevos y emocionantes. Estos clientes están investigando modelos base, como TII Falcon, Stable Diffusion XL o GPT-3.5 de OpenAI, como los motores que impulsan la innovación en IA generativa.

Los modelos base son una clase de modelos de IA generativa capaces de comprender y generar contenido similar al humano, gracias a las grandes cantidades de datos no estructurados en los que han sido entrenados. Estos modelos han revolucionado varias tareas de visión por computadora (CV) y procesamiento de lenguaje natural (NLP), incluyendo generación de imágenes, traducción y respuesta a preguntas. Sirven como bloques de construcción para muchas aplicaciones de IA y se han convertido en un componente crucial en el desarrollo de sistemas inteligentes avanzados.

Sin embargo, la implementación de modelos base puede presentar desafíos significativos, especialmente en términos de costos y requisitos de recursos. Estos modelos son conocidos por su tamaño, que a menudo varía desde cientos de millones hasta miles de millones de parámetros. Su gran tamaño requiere de recursos computacionales extensos, incluyendo hardware potente y una capacidad de memoria significativa. De hecho, desplegar modelos base generalmente requiere al menos una (a menudo más) GPU para manejar eficientemente la carga computacional. Por ejemplo, el modelo TII Falcon-40B Instruct requiere al menos una instancia ml.g5.12xlarge para poder ser cargado en memoria correctamente, pero funciona mejor con instancias más grandes. Como resultado, el retorno de inversión (ROI) de implementar y mantener estos modelos puede ser demasiado bajo para demostrar valor empresarial, especialmente durante los ciclos de desarrollo o para cargas de trabajo puntuales. Esto se debe a los costos de operación de tener instancias con GPU durante largas sesiones, potencialmente las 24 horas del día, los 7 días de la semana.

A principios de este año, anunciamos Amazon Bedrock, una API sin servidor para acceder a modelos base de Amazon y nuestros socios de IA generativa. Aunque actualmente se encuentra en Vista previa privada, su API sin servidor le permite utilizar modelos base de Amazon, Anthropic, Stability AI y AI21, sin tener que implementar ningún punto de conexión usted mismo. Sin embargo, los modelos de código abierto de comunidades como Hugging Face han estado creciendo mucho y no todos ellos han sido puestos a disposición a través de Amazon Bedrock.

• Data Morph Avanzando más allá de la docena de Datasaurus
• Aprendizaje de representación en grafos y redes
• Construyendo un Asistente con ChatGPT y Voz habilitada utilizando React y Express

En esta publicación, nos enfocamos en estas situaciones y resolvemos el problema de los altos costos al implementar grandes modelos base en puntos de conexión asíncronos de Amazon SageMaker desde Amazon SageMaker JumpStart. Esto puede ayudar a reducir los costos de la arquitectura, permitiendo que el punto de conexión se ejecute solo cuando haya solicitudes en la cola y por un corto tiempo de vida, mientras se reduce a cero cuando no haya solicitudes esperando ser atendidas. Esto suena genial para muchos casos de uso; sin embargo, un punto de conexión que se ha reducido a cero agregará un tiempo de inicio en frío antes de poder servir inferencias.

Descripción general de la solución

El siguiente diagrama ilustra la arquitectura de nuestra solución.

La arquitectura que implementamos es muy sencilla:

• La interfaz de usuario es un cuaderno, que puede ser reemplazado por una interfaz de usuario web construida en Streamlit u otra tecnología similar. En nuestro caso, el cuaderno es un cuaderno de Amazon SageMaker Studio, que se ejecuta en una instancia ml.m5.large con el kernel PyTorch 2.0 Python 3.10 CPU.
• El cuaderno hace consultas al punto de conexión de tres maneras: a través del SDK de Python de SageMaker, el SDK de AWS para Python (Boto3) y LangChain.
• El punto de conexión se ejecuta de forma asíncrona en SageMaker, y en el punto de conexión, implementamos el modelo Falcon-40B Instruct. Actualmente es el estado del arte en términos de modelos instructivos y está disponible en SageMaker JumpStart. Una sola llamada a la API nos permite implementar el modelo en el punto de conexión.

¿Qué es la inferencia asíncrona de SageMaker?

La inferencia asíncrona de SageMaker es una de las cuatro opciones de implementación en SageMaker, junto con puntos de conexión en tiempo real, inferencia por lotes e inferencia sin servidor. Para obtener más información sobre las diferentes opciones de implementación, consulte Implementar modelos para inferencia.

La inferencia asíncrona de SageMaker encola las solicitudes entrantes y las procesa de manera asíncrona, lo que hace que esta opción sea ideal para solicitudes con tamaños de carga útil grandes de hasta 1 GB, tiempos de procesamiento largos y requisitos de latencia casi en tiempo real. Sin embargo, la principal ventaja que ofrece al tratar con grandes modelos base, especialmente durante una prueba de concepto (POC) o durante el desarrollo, es la capacidad de configurar la inferencia asíncrona para escalar a un número de instancias igual a cero cuando no hay solicitudes que procesar, lo que permite ahorrar costos. Para obtener más información sobre la inferencia asíncrona de SageMaker, consulte Inferencia asíncrona. El siguiente diagrama ilustra esta arquitectura.

Para implementar un punto de enlace de inferencia asíncrona, necesitas crear un objeto AsyncInferenceConfig. Si creas AsyncInferenceConfig sin especificar sus argumentos, el S3OutputPath predeterminado será s3://sagemaker-{REGION}-{ACCOUNTID}/async-endpoint-outputs/{UNIQUE-JOB-NAME} y el S3FailurePath será s3://sagemaker-{REGION}-{ACCOUNTID}/async-endpoint-failures/{UNIQUE-JOB-NAME}.

¿Qué es SageMaker JumpStart?

Nuestro modelo proviene de SageMaker JumpStart, una característica de SageMaker que acelera el viaje del aprendizaje automático (ML) ofreciendo modelos pre-entrenados, plantillas de solución y cuadernos de ejemplo. Proporciona acceso a una amplia gama de modelos pre-entrenados para diferentes tipos de problemas, lo que te permite comenzar tus tareas de ML con una base sólida. SageMaker JumpStart también ofrece plantillas de solución para casos de uso comunes y cuadernos de ejemplo para aprender. Con SageMaker JumpStart, puedes reducir el tiempo y el esfuerzo necesarios para comenzar tus proyectos de ML con lanzamientos de soluciones con un solo clic y recursos completos para una experiencia práctica de ML.

La siguiente captura de pantalla muestra un ejemplo de solo algunos de los modelos disponibles en la interfaz de usuario de SageMaker JumpStart.

Implementar el modelo

Nuestro primer paso es implementar el modelo en SageMaker. Para hacerlo, podemos usar la interfaz de usuario de SageMaker JumpStart o el SDK de Python de SageMaker, que proporciona una API que podemos usar para implementar el modelo en el punto de enlace asíncrono:

%%time
from sagemaker.jumpstart.model import JumpStartModel, AsyncInferenceConfig
from sagemaker.serializers import JSONSerializer
from sagemaker.deserializers import JSONDeserializer

model_id, model_version = "huggingface-llm-falcon-40b-instruct-bf16", "*"
my_model = JumpStartModel(model_id=model_id)
predictor = my_model.deploy(
    initial_instance_count=0,
    instance_type="ml.g5.12xlarge",
    async_inference_config=AsyncInferenceConfig()
)

Esta llamada puede tardar aproximadamente 10 minutos en completarse. Durante este tiempo, se inicia el punto de enlace, se descargan el contenedor junto con los artefactos del modelo en el punto de enlace, se carga la configuración del modelo desde SageMaker JumpStart y luego se expone el punto de enlace asíncrono a través de un punto de enlace DNS. Para asegurarnos de que nuestro punto de enlace pueda reducirse a cero, necesitamos configurar el escalado automático en el punto de enlace asíncrono utilizando Application Auto Scaling. Primero debes registrar tu variante de punto de enlace con Application Auto Scaling, definir una política de escalado y luego aplicar la política de escalado. En esta configuración, usamos una métrica personalizada utilizando CustomizedMetricSpecification, llamada ApproximateBacklogSizePerInstance, como se muestra en el siguiente código. Para obtener una lista detallada de las métricas de Amazon CloudWatch disponibles con tu punto de enlace de inferencia asíncrona, consulta la documentación sobre Monitorización con CloudWatch.

import boto3

client = boto3.client("application-autoscaling")
resource_id = "endpoint/" + my_model.endpoint_name + "/variant/" + "AllTraffic"

# Configurar el escalado automático en el punto de enlace asíncrono hasta cero instancias
response = client.register_scalable_target(
    ServiceNamespace="sagemaker",
    ResourceId=resource_id,
    ScalableDimension="sagemaker:variant:DesiredInstanceCount",
    MinCapacity=0, # Número mínimo de instancias al que queremos reducir el escalado - reducir a 0 para dejar de incurrir en costos
    MaxCapacity=1, # Número máximo de instancias al que queremos aumentar el escalado - aumentar a 1 máximo es suficiente para desarrollo
)

response = client.put_scaling_policy(
    PolicyName="Invocations-ScalingPolicy",
    ServiceNamespace="sagemaker",  # El espacio de nombres del servicio de AWS que proporciona el recurso.
    ResourceId=resource_id,  # Nombre del punto de enlace
    ScalableDimension="sagemaker:variant:DesiredInstanceCount",  # SageMaker solo admite conteo de instancias
    PolicyType="TargetTrackingScaling",  # 'StepScaling'|'TargetTrackingScaling'
    TargetTrackingScalingPolicyConfiguration={
        "TargetValue": 5.0,  # El valor objetivo para la métrica. - aquí la métrica es - SageMakerVariantInvocationsPerInstance
        "CustomizedMetricSpecification": {
            "MetricName": "ApproximateBacklogSizePerInstance",
            "Namespace": "AWS/SageMaker",
            "Dimensions": [{"Name": "EndpointName", "Value": my_model.endpoint_name}],
            "Statistic": "Average",
        },
        "ScaleInCooldown": 600,  # La cantidad de tiempo, en segundos, después de que se completa una actividad de reducción de escala antes de que pueda comenzar otra actividad de reducción de escala.
        "ScaleOutCooldown": 300,  # ScaleOutCooldown - La cantidad de tiempo, en segundos, después de que se completa una actividad de aumento de escala antes de que pueda comenzar otra actividad de aumento de escala.
        # 'DisableScaleIn': True|False - indica si está deshabilitada la reducción de escala por la política de seguimiento de objetivo.
        # Si el valor es verdadero, la reducción de escala está deshabilitada y la política de seguimiento de objetivo no eliminará capacidad del recurso escalable.
    },
)

Puedes verificar que esta política se haya configurado correctamente navegando a la consola de SageMaker, seleccionando Endpoints bajo Inference en el panel de navegación, y buscando el endpoint que acabamos de implementar.

Invocar el endpoint asincrónico

Para invocar el endpoint, debes colocar el payload de la solicitud en Amazon Simple Storage Service (Amazon S3) y proporcionar un puntero a este payload como parte de la solicitud InvokeEndpointAsync. Al invocar, SageMaker encola la solicitud para su procesamiento y devuelve un identificador y una ubicación de salida como respuesta. Una vez procesado, SageMaker coloca el resultado en la ubicación de Amazon S3. Opcionalmente, puedes elegir recibir notificaciones de éxito o error con Amazon Simple Notification Service (Amazon SNS).

SageMaker Python SDK

Una vez que se complete la implementación, se devolverá un objeto AsyncPredictor. Para realizar inferencia asincrónica, debes cargar los datos en Amazon S3 y utilizar el método predict_async() con la URI de S3 como entrada. Devolverá un objeto AsyncInferenceResponse, y puedes verificar el resultado utilizando el método get_response().

Alternativamente, si deseas verificar un resultado periódicamente y devolverlo cuando esté generado, utiliza el método predict(). Utilizamos este segundo método en el siguiente código:

import time

# Invocar el endpoint asincrónico con el SageMaker Python SDK
def query_endpoint(payload):
    """Consultar el endpoint e imprimir la respuesta"""
    response = predictor.predict_async(
        data=payload,
        input_path="s3://{}/{}".format(bucket, prefix),
    )
    while True:
        try:
            response = response.get_result()
            break
        except:
            print("La inferencia no está lista...")
            time.sleep(5)
    print(f"\033[1m Input:\033[0m {payload['inputs']}")
    print(f"\033[1m Output:\033[0m {response[0]['generated_text']}")
    
query_endpoint(payload)

Boto3

Ahora vamos a explorar el método invoke_endpoint_async del cliente sagemaker-runtime de Boto3. Permite a los desarrolladores invocar de forma asíncrona un endpoint de SageMaker, proporcionando un token para el seguimiento del progreso y la recuperación de la respuesta más tarde. Boto3 no ofrece una forma de esperar a que se complete la inferencia asíncrona como la operación get_result() del SageMaker Python SDK. Por lo tanto, aprovechamos el hecho de que Boto3 almacenará la salida de la inferencia en Amazon S3 en response["OutputLocation"]. Podemos utilizar la siguiente función para esperar a que se escriba el archivo de inferencia en Amazon S3:

import json
import time
import boto3
from botocore.exceptions import ClientError

s3_client = boto3.client("s3")

# Esperar hasta que se genere la predicción
def wait_inference_file(bucket, prefix):
    while True:
        try:
            response = s3_client.get_object(Bucket=bucket, Key=prefix)
            break
        except ClientError as ex:
            if ex.response['Error']['Code'] == 'NoSuchKey':
                print("Esperando a que se genere el archivo...")
                time.sleep(5)
                next
            else:
                raise
        except Exception as e:
            print(e.__dict__)
            raise
    return response

Con esta función, ahora podemos consultar el endpoint:

# Invocar el endpoint asincrónico con el SDK de Boto3
import boto3

sagemaker_client = boto3.client("sagemaker-runtime")

# Función para consultar el endpoint
def query_endpoint_boto3(payload):
    """Consultar el endpoint e imprimir la respuesta"""
    response = sagemaker_client.invoke_endpoint_async(
        EndpointName=my_model.endpoint_name,
        InputLocation="s3://{}/{}".format(bucket, prefix),
        ContentType="application/json",
        Accept="application/json"
    )
    output_url = response["OutputLocation"]
    output_prefix = "/".join(output_url.split("/")[3:])
    # Leer los bytes del archivo de S3 en output_url con Boto3
    output = wait_inference_file(bucket, output_prefix)
    output = json.loads(output['Body'].read())[0]['generated_text']
    # Emitir la salida
    print(f"\033[1m Input:\033[0m {payload['inputs']}")
    print(f"\033[1m Output:\033[0m {output}")

query_endpoint_boto3(payload)

LangChain

LangChain es un marco de código abierto lanzado en octubre de 2022 por Harrison Chase. Simplifica el desarrollo de aplicaciones que utilizan modelos de lenguaje grandes (LLMs) al proporcionar integraciones con varios sistemas y fuentes de datos. LangChain permite el análisis de documentos, la creación de resúmenes, la creación de chatbots, el análisis de código y más. Ha ganado popularidad, con contribuciones de cientos de desarrolladores y financiamiento significativo de empresas de capital de riesgo. LangChain permite la conexión de LLMs con fuentes externas, lo que hace posible crear aplicaciones dinámicas y sensibles a los datos. Ofrece bibliotecas, APIs y documentación para agilizar el proceso de desarrollo.

LangChain proporciona bibliotecas y ejemplos para utilizar puntos finales de SageMaker con su marco, lo que facilita el uso de modelos de ML alojados en SageMaker como el “cerebro” de la cadena. Para obtener más información sobre cómo LangChain se integra con SageMaker, consulte el punto final de SageMaker en la documentación de LangChain.

Una de las limitaciones de la implementación actual de LangChain es que no admite puntos finales asíncronos de forma nativa. Para utilizar un punto final asíncrono en LangChain, debemos definir una nueva clase, SagemakerAsyncEndpoint, que extienda la clase SagemakerEndpoint ya disponible en LangChain. Además, proporcionamos la siguiente información:

• El bucket y prefijo de S3 donde la inferencia asíncrona almacenará las entradas (y salidas)
• Un número máximo de segundos para esperar antes de agotar el tiempo de espera
• Una función _call() actualizada para consultar el punto final con invoke_endpoint_async() en lugar de invoke_endpoint()
• Una forma de despertar el punto final asíncrono si está en inicio en frío (escalado a cero)

Para revisar el SagemakerAsyncEndpoint recién creado, puedes consultar el archivo sagemaker_async_endpoint.py disponible en GitHub.

from typing import Dict
from langchain import PromptTemplate
from langchain.llms.sagemaker_endpoint import LLMContentHandler
from langchain.chains import LLMChain
from sagemaker_async_endpoint import SagemakerAsyncEndpoint

class ContentHandler(LLMContentHandler):
    content_type:str = "application/json"
    accepts:str = "application/json"
    len_prompt:int = 0

    def transform_input(self, prompt: str, model_kwargs: Dict) -> bytes:
        self.len_prompt = len(prompt)
        input_str = json.dumps({"inputs": prompt, "parameters": {"max_new_tokens": 100, "do_sample": False, "repetition_penalty": 1.1}})
        return input_str.encode('utf-8')

    def transform_output(self, output: bytes) -> str:
        response_json = output.read()
        res = json.loads(response_json)
        ans = res[0]['generated_text']
        return ans

chain = LLMChain(
    llm=SagemakerAsyncEndpoint(
        input_bucket=bucket,
        input_prefix=prefix,
        endpoint_name=my_model.endpoint_name,
        region_name=sagemaker.Session().boto_region_name,
        content_handler=ContentHandler(),
    ),
    prompt=PromptTemplate(
        input_variables=["query"],
        template="{query}",
    ),
)

print(chain.run(payload['inputs']))

Limpieza

Cuando hayas terminado de probar la generación de inferencias desde el punto final, recuerda eliminar el punto final para evitar incurrir en cargos adicionales:

predictor.delete_endpoint()

Conclusión

Cuando se implementan modelos de base grandes como TII Falcon, optimizar los costos es crucial. Estos modelos requieren hardware potente y una capacidad de memoria sustancial, lo que conlleva altos costos de infraestructura. La inferencia asíncrona de SageMaker, una opción de implementación que procesa solicitudes de forma asíncrona, reduce los gastos al escalar el número de instancias a cero cuando no hay solicitudes pendientes. En esta publicación, demostramos cómo implementar modelos de base de SageMaker JumpStart en puntos finales asíncronos de SageMaker. Proporcionamos ejemplos de código utilizando el SDK de Python de SageMaker, Boto3 y LangChain para ilustrar diferentes métodos para invocar puntos finales asíncronos y obtener resultados. Estas técnicas permiten a los desarrolladores e investigadores optimizar los costos mientras utilizan las capacidades de los modelos de base para sistemas avanzados de comprensión del lenguaje.

Para obtener más información sobre la inferencia asíncrona y SageMaker JumpStart, consulta las siguientes publicaciones:

• Crear rápidamente aplicaciones de IA generativas de alta precisión en datos empresariales utilizando Amazon Kendra, LangChain y grandes modelos de lenguaje
• Ejecutar inferencias de visión por computadora en videos grandes con puntos finales asíncronos de Amazon SageMaker

We will continue to update Zepes; if you have any questions or suggestions, please contact us!