Maximiza el rendimiento estable de la difusión y reduce los costos de inferencia con AWS Inferentia2

Maximiza el rendimiento y reduce los costos con AWS Inferentia2

Los modelos de IA generativa han experimentado un crecimiento rápido en los últimos meses debido a sus impresionantes capacidades para crear texto, imágenes, código y audio realistas. Entre estos modelos, destacan los modelos de Difusión Estable por su fuerza única en la creación de imágenes de alta calidad basadas en indicaciones de texto. La Difusión Estable puede generar una amplia variedad de imágenes de alta calidad, incluyendo retratos realistas, paisajes e incluso arte abstracto. Y, al igual que otros modelos de IA generativa, los modelos de Difusión Estable requieren una potente computación para proporcionar inferencia de baja latencia.

En esta publicación, mostramos cómo puedes ejecutar modelos de Difusión Estable y lograr un alto rendimiento al menor costo en Amazon Elastic Compute Cloud (Amazon EC2) utilizando instancias Inf2 de Amazon EC2 alimentadas por AWS Inferentia2. Analizamos la arquitectura de un modelo de Difusión Estable y recorremos los pasos para compilar un modelo de Difusión Estable utilizando AWS Neuron y desplegarlo en una instancia Inf2. También discutimos las optimizaciones que el SDK Neuron realiza automáticamente para mejorar el rendimiento. Puedes ejecutar tanto las versiones 2.1 como 1.5 de Difusión Estable de manera rentable en AWS Inferentia2. Por último, mostramos cómo puedes implementar un modelo de Difusión Estable en una instancia Inf2 con Amazon SageMaker.

El tamaño del modelo de Difusión Estable 2.1 en punto flotante 32 (FP32) es de 5 GB y de 2.5 GB en bfoat16 (BF16). Una instancia inf2.xlarge tiene un acelerador AWS Inferentia2 con 32 GB de memoria HBM. El modelo de Difusión Estable 2.1 cabe en una sola instancia inf2.xlarge. La Difusión Estable es un modelo de texto a imagen que puedes utilizar para crear imágenes de diferentes estilos y contenido simplemente proporcionando una indicación de texto como entrada. Para obtener más información sobre la arquitectura del modelo de Difusión Estable, consulta “Crear imágenes de alta calidad con modelos de Difusión Estable y desplegarlas de manera rentable con Amazon SageMaker”.

Cómo el SDK Neuron optimiza el rendimiento de Difusión Estable

Antes de poder implementar el modelo de Difusión Estable 2.1 en instancias de AWS Inferentia2, debemos compilar los componentes del modelo utilizando el SDK Neuron. El SDK Neuron, que incluye un compilador, tiempo de ejecución y herramientas de aprendizaje profundo, compila y optimiza automáticamente los modelos de aprendizaje profundo para que se ejecuten de manera eficiente en las instancias Inf2 y aprovechen todo el rendimiento del acelerador AWS Inferentia2. Tenemos ejemplos disponibles para el modelo de Difusión Estable 2.1 en el repositorio de GitHub. Este cuaderno presenta un ejemplo completo de cómo compilar un modelo de Difusión Estable, guardar los modelos compilados de Neuron y cargarlos en el tiempo de ejecución para inferencia.

Utilizamos StableDiffusionPipeline de la biblioteca diffusers de Hugging Face para cargar y compilar el modelo. Luego compilamos todos los componentes del modelo para Neuron utilizando torch_neuronx.trace() y guardamos el modelo optimizado como TorchScript. Los procesos de compilación pueden requerir una cantidad significativa de memoria RAM. Para evitar esto, antes de trazar cada modelo, creamos una copia profunda de la parte de la canalización que se está trazando. Después, eliminamos el objeto de la canalización de la memoria utilizando del pipe. Esta técnica es especialmente útil al compilar en instancias con poca RAM.

Además, también realizamos optimizaciones a los modelos de Difusión Estable. UNet es la parte más intensiva computacionalmente de la inferencia. El componente UNet opera en tensores de entrada que tienen un tamaño de lote de dos, generando un tensor de salida correspondiente también con un tamaño de lote de dos, para producir una sola imagen. Los elementos dentro de estos lotes son completamente independientes entre sí. Podemos aprovechar este comportamiento para obtener una latencia óptima ejecutando un lote en cada núcleo de Neuron. Compilamos UNet para un lote (utilizando tensores de entrada con un lote) y luego utilizamos la API torch_neuronx.DataParallel para cargar este modelo de un solo lote en cada núcleo. La salida de esta API es un módulo de dos lotes sin problemas: podemos pasar al UNet las entradas de dos lotes y se devuelve una salida de dos lotes, pero internamente, los dos modelos de un solo lote se están ejecutando en los dos núcleos de Neuron. Esta estrategia optimiza la utilización de recursos y reduce la latencia.

Compilar y desplegar un modelo de Difusión Estable en una instancia EC2 Inf2

Para compilar y desplegar el modelo de Difusión Estable en una instancia EC2 Inf2, inicia sesión en la Consola de administración de AWS y crea una instancia inf2.8xlarge. Ten en cuenta que solo se requiere una instancia inf2.8xlarge para la compilación del modelo, ya que la compilación requiere una memoria principal más alta. El modelo de Difusión Estable se puede alojar en una instancia inf2.xlarge. Puedes encontrar la última AMI con las bibliotecas de Neuron utilizando el siguiente comando de la Interfaz de línea de comandos de AWS (AWS CLI):

aws ec2 describe-images --region us-east-1 --owners amazon \
--filters 'Name=name,Values=Deep Learning AMI Neuron PyTorch 1.13.? (Amazon Linux 2) ????????' 'Name=state,Values=available' \
--query 'reverse(sort_by(Images, &CreationDate))[:1].ImageId' \
--output text

Para este ejemplo, creamos una instacia EC2 utilizando la AMI de Deep Learning Neuron PyTorch 1.13 (Ubuntu 20.04). Luego puedes crear un entorno de laboratorio JupyterLab conectándote a la instancia y ejecutando los siguientes pasos:

ejecuta source /opt/aws_neuron_venv_pytorch/bin/activate
pip install jupyterlab
jupyter-lab

Un cuaderno con todos los pasos para compilar y alojar el modelo se encuentra en GitHub.

Echemos un vistazo a los pasos de compilación para uno de los bloques de codificador de texto. Otros bloques que forman parte del pipeline Stable Diffusion se pueden compilar de manera similar.

El primer paso es cargar el modelo pre-entrenado de Hugging Face. El método StableDiffusionPipeline.from_pretrained carga el modelo pre-entrenado en nuestro objeto de pipeline, pipe. Luego creamos una copia profunda del codificador de texto de nuestro pipeline, clonándolo efectivamente. El comando del pipe se usa para eliminar el objeto de pipeline original, liberando la memoria que consumía. Aquí, estamos cuantizando el modelo a pesos BF16:

model_id = "stabilityai/stable-diffusion-2-1-base"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
text_encoder = copy.deepcopy(pipe.text_encoder)
del pipe

Este paso implica envolver nuestro codificador de texto con el envoltorio NeuronTextEncoder. La salida de un módulo de codificador de texto compilado será de tipo dict. Lo convertimos a un tipo list usando este envoltorio:

text_encoder = NeuronTextEncoder(text_encoder)

Inicializamos el tensor de PyTorch emb con algunos valores. El tensor emb se utiliza como entrada de ejemplo para la función torch_neuronx.trace. Esta función traza nuestro codificador de texto y lo compila en un formato optimizado para Neuron. La ruta de directorio para el modelo compilado se construye uniéndo COMPILER_WORKDIR_ROOT con el subdirectorio text_encoder:

emb = torch.tensor([...])
text_encoder_neuron = torch_neuronx.trace(
       text_encoder.neuron_text_encoder,
       emb,
       compiler_workdir=os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder'),
       )

El codificador de texto compilado se guarda utilizando torch.jit.save. Se guarda con el nombre de archivo model.pt en el directorio text_encoder del espacio de trabajo de nuestro compilador:

text_encoder_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder/model.pt')
torch.jit.save(text_encoder_neuron, text_encoder_filename)

El cuaderno incluye pasos similares para compilar otros componentes del modelo: UNet, VAE decoder y VAE post_quant_conv. Después de compilar todos los modelos, puedes cargar y ejecutar el modelo siguiendo estos pasos:

  1. Define las rutas para los modelos compilados.
  2. Carga un modelo pre-entrenado de StableDiffusionPipeline, con su configuración especificada para usar el tipo de datos bfloat16.
  3. Carga el modelo UNet en dos núcleos de Neuron utilizando la API torch_neuronx.DataParallel. Esto permite realizar inferencia paralela de datos, lo que puede acelerar significativamente el rendimiento del modelo.
  4. Carga las partes restantes del modelo (text_encoder, decoder y post_quant_conv) en un solo núcleo de Neuron.

Luego puedes ejecutar el pipeline proporcionando texto de entrada como inicios. A continuación se muestran algunas imágenes generadas por el modelo para los inicios:

  • Retrato de renaud sechan, tinta y tinta china, dibujos de líneas intrincadas, de craig mullins, ruan jia, kentaro miura, greg rutkowski, loundraw

  • Retrato de un antiguo minero de carbón en el siglo XIX, hermosa pintura, con un retrato facial altamente detallado realizado por Greg Rutkowski

  • Un castillo en medio de un bosque

Hospedar Stable Diffusion 2.1 en AWS Inferentia2 y SageMaker

Hospedar modelos de Stable Diffusion con SageMaker también requiere compilación con Neuron SDK. Puede completar la compilación con anticipación o durante la ejecución utilizando contenedores de Inferencia de Modelos Grandes (LMI, por sus siglas en inglés). La compilación anticipada permite tiempos de carga de modelo más rápidos y es la opción preferida.

Los contenedores LMI de SageMaker proporcionan dos formas de implementar el modelo:

  • Una opción sin código en la que solo proporcionamos un archivo serving.properties con las configuraciones requeridas
  • Utilizar su propio script de inferencia

Examinamos ambas soluciones y repasamos las configuraciones y el script de inferencia (model.py). En esta publicación, demostramos la implementación utilizando un modelo precompilado almacenado en un bucket de Amazon Simple Storage Service (Amazon S3). Puede usar este modelo precompilado para sus implementaciones.

Configurar el modelo con un script proporcionado

En esta sección, mostramos cómo configurar el contenedor LMI para hospedar los modelos de Stable Diffusion. El cuaderno SD2.1 está disponible en GitHub. El primer paso es crear el paquete de configuración del modelo según la siguiente estructura de directorios. Nuestro objetivo es utilizar las configuraciones de modelo mínimas necesarias para hospedar el modelo. La estructura de directorios necesaria es la siguiente:

<directorio-raíz-de-configuración> / 
    ├── serving.properties
    │   
    └── model.py [OPCIONAL]

A continuación, creamos el archivo serving.properties con los siguientes parámetros:

%%writefile code_sd/serving.properties
engine=Python
option.entryPoint=djl_python.transformers-neuronx
option.use_stable_diffusion=True
option.model_id=s3url
option.tensor_parallel_degree=2
option.dtype=bf16

Los parámetros especifican lo siguiente:

  • option.model_id – Los contenedores LMI utilizan s5cmd para cargar el modelo desde la ubicación de S3 y, por lo tanto, debemos especificar la ubicación de nuestros pesos compilados.
  • option.entryPoint – Para utilizar los controladores incorporados, especificamos la clase transformers-neuronx. Si tiene un script de inferencia personalizado, debe proporcionarlo en su lugar.
  • option.dtype – Esto especifica cargar los pesos en un tamaño específico. Para esta publicación, utilizamos BF16, lo que reduce aún más nuestros requisitos de memoria y disminuye nuestra latencia debido a eso.
  • option.tensor_parallel_degree – Este parámetro especifica el número de aceleradores que utilizamos para este modelo. El acelerador de chip AWS Inferentia2 tiene dos núcleos de Neuron, por lo que especificar un valor de 2 significa que utilizamos un acelerador (dos núcleos). Esto significa que ahora podemos crear varios trabajadores para aumentar el rendimiento del punto de enlace.
  • option.engine – Esto se establece en Python para indicar que no utilizaremos otros compiladores como DeepSpeed o Faster Transformer para este hospedaje.

Utilizar su propio script

Si desea utilizar su propio script de inferencia personalizado, debe eliminar option.entryPoint de serving.properties. El contenedor LMI, en ese caso, buscará un archivo model.py en la misma ubicación que serving.properties y lo utilizará para ejecutar la inferencia.

Crear tu propio script de inferencia (model.py)

Crear tu propio script de inferencia es relativamente sencillo utilizando el contenedor LMI. El contenedor requiere que tu archivo model.py tenga una implementación del siguiente método:

def handle(inputs: Input) que devuelve un objeto de tipo Outputs

Examinemos algunas de las áreas críticas del cuaderno adjunto, que demuestra la función de traer tu propio script.

Reemplaza el módulo cross_attention con la versión optimizada:

# Reemplazar el módulo de atención cruzada original con un módulo de atención cruzada personalizado para obtener un mejor rendimiento
    CrossAttention.get_attention_scores = get_attention_scores
Cargar los pesos compilados para lo siguiente
text_encoder_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder.pt')
decoder_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'vae_decoder.pt')
unet_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'unet.pt')
post_quant_conv_filename =. os.path.join(COMPILER_WORKDIR_ROOT, 'vae_post_quant_conv.pt')

Estos son los nombres de los archivos de pesos compilados que utilizamos al crear las compilaciones. Siéntete libre de cambiar los nombres de los archivos, pero asegúrate de que los nombres de tus archivos de pesos coincidan con lo que especifiques aquí.

Luego debemos cargarlos utilizando el SDK de Neuron y establecerlos en los pesos reales del modelo. Al cargar los pesos optimizados de UNet, ten en cuenta que también especificamos el número de núcleos de Neuron que necesitamos cargar en ellos. Aquí, cargamos en un solo acelerador con dos núcleos:

# Cargar el UNet compilado en dos núcleos de Neuron.
    pipe.unet = NeuronUNet(UNetWrap(pipe.unet))
    logging.info(f"Cargando modelo: unet: creado")
    device_ids = [idx for idx in range(tensor_parallel_degree)]
   
    pipe.unet.unetwrap = torch_neuronx.DataParallel(torch.jit.load(unet_filename), device_ids, set_dynamic_batching=False)
   
 
    # Cargar otros modelos compilados en un solo núcleo de Neuron.
 
    # - cargar codificadores
    pipe.text_encoder = NeuronTextEncoder(pipe.text_encoder)
    clip_compiled = torch.jit.load(text_encoder_filename)
    pipe.text_encoder.neuron_text_encoder = clip_compiled
    #- cargar decodificadores
    pipe.vae.decoder = torch.jit.load(decoder_filename)
    pipe.vae.post_quant_conv = torch.jit.load(post_quant_conv_filename)

Ejecutar la inferencia con un indicador invoca el objeto pipe para generar una imagen.

Crear el punto de enlace de SageMaker

Utilizamos las API de Boto3 para crear un punto de enlace de SageMaker. Completa los siguientes pasos:

  1. Crea el archivo tarball solo con los archivos de servicio y, opcionalmente, model.py, y súbalo a Amazon S3.
  2. Crea el modelo utilizando el contenedor de imagen y el tarball del modelo cargado anteriormente.
  3. Crea la configuración del punto de enlace utilizando los siguientes parámetros clave:
    1. Utiliza una instancia ml.inf2.xlarge.
    2. Establece ContainerStartupHealthCheckTimeoutInSeconds en 240 para asegurarte de que la verificación de salud comience después de que se implemente el modelo.
    3. Establece VolumeInGB en un valor mayor para que pueda utilizarse para cargar los pesos del modelo que tienen un tamaño de 32 GB.

Crear un modelo de SageMaker

Después de crear el archivo model.tar.gz y subirlo a Amazon S3, necesitamos crear un modelo de SageMaker. Utilizamos el contenedor LMI y el artefacto del modelo del paso anterior para crear el modelo de SageMaker. SageMaker nos permite personalizar e inyectar varias variables de entorno. Para este flujo de trabajo, podemos dejar todo como predeterminado. Observa el siguiente código:

inference_image_uri = (
    f"763104351884.dkr.ecr.{region}.amazonaws.com/djl-inference:0 djl-serving-inf2"
)

Crea el objeto de modelo, que esencialmente crea un contenedor de bloqueo que se carga en la instancia y se utiliza para inferir:

model_name = name_from_base(f"inf2-sd")
create_model_response = boto3_sm_client.create_model(
    ModelName=model_name,
    ExecutionRoleArn=role,
    PrimaryContainer={"Image": inference_image_uri, "ModelDataUrl": s3_code_artifact},
)

Crear un punto de enlace de SageMaker

En esta demostración, utilizamos una instancia ml.inf2.xlarge. Necesitamos establecer los parámetros VolumeSizeInGB para proporcionar el espacio en disco necesario para cargar el modelo y los pesos. Este parámetro se aplica a las instancias que admiten la conexión de volumen de Amazon Elastic Block Store (Amazon EBS). Podemos dejar el tiempo de espera de descarga del modelo y la comprobación de salud del inicio del contenedor con un valor más alto, lo que dará tiempo suficiente para que el contenedor extraiga los pesos de Amazon S3 y los cargue en los aceleradores AWS Inferentia2. Para obtener más detalles, consulte CreateEndpointConfig.

endpoint_config_response = boto3_sm_client.create_endpoint_config(

EndpointConfigName=endpoint_config_name,
    ProductionVariants=[
        {
            "VariantName": "variant1",
            "ModelName": model_name,
            "InstanceType": "ml.inf2.xlarge", # - 
            "InitialInstanceCount": 1,
            "ContainerStartupHealthCheckTimeoutInSeconds": 360, 
            "VolumeSizeInGB": 400
        },
    ],
)

Por último, creamos un punto de enlace de SageMaker:

create_endpoint_response = boto3_sm_client.create_endpoint(
    EndpointName=f"{endpoint_name}", EndpointConfigName=endpoint_config_name
)

Invocar el punto de enlace del modelo

Este es un modelo generativo, por lo que pasamos la indicación que el modelo utiliza para generar la imagen. La carga útil es de tipo JSON:

response_model = boto3_sm_run_client.invoke_endpoint(

EndpointName=endpoint_name,
    Body=json.dumps(
        {
            "prompt": "Paisaje de montaña", 
            "parameters": {} # 
        }
    ), 
    ContentType="application/json",
)

Benchmarkear el modelo Stable Diffusion en Inf2

Ejecutamos algunas pruebas para evaluar el rendimiento del modelo Stable Diffusion con el tipo de dato BF 16 en Inf2, y pudimos obtener números de latencia que rivalizan o superan a algunos de los otros aceleradores para Stable Diffusion. Esto, junto con el menor costo de las chips AWS Inferentia2, lo convierte en una propuesta extremadamente valiosa.

Los siguientes números son del modelo Stable Diffusion implementado en una instancia inf2.xl. Para obtener más información sobre los costos, consulte Amazon EC2 Inf2 Instances.

Modelo Stable Diffusion 1.5 Stable Diffusion 1.5 Stable Diffusion 1.5 Stable Diffusion 1.5 Resolución 512×512 768×768 512×512 768×768 Tipo de dato bf16 bf16 bf16 bf16 Iteraciones 50 50 30 30 Latencia P95 (ms) 2,427.4 8,235.9 1,456.5 4,941.6 Costo bajo demanda de Inf2.xl por hora $0.76 $0.76 $0.76 $0.76 Inf2.xl (Costo por imagen) $0.0005125 $0.0017387 $0.0003075 $0.0010432
Stable Diffusion 2.1 Stable Diffusion 2.1 Stable Diffusion 2.1 Stable Diffusion 2.1 512×512 768×768 512×512 768×768 bf16 bf16 bf16 bf16 50 50 30 30 1,976.9 6,836.3 1,186.2 4,101.8 $0.76 $0.76 $0.76 $0.76 $0.0004174 $0.0014432 $0.0002504 $0.0008659

Conclusión

En esta publicación, nos sumergimos en la compilación, optimización y implementación del modelo Stable Diffusion 2.1 utilizando instancias Inf2. También demostramos la implementación de modelos Stable Diffusion utilizando SageMaker. Las instancias Inf2 también ofrecen un excelente rendimiento en términos de precio para Stable Diffusion 1.5. Para obtener más información sobre por qué las instancias Inf2 son ideales para la inteligencia artificial generativa y los modelos de lenguaje grandes, consulte “Amazon EC2 Inf2 Instances for Low-Cost, High-Performance Generative AI Inference are Now Generally Available”. Para conocer más detalles de rendimiento, consulte “Inf2 Performance”. Eche un vistazo a ejemplos adicionales en el repositorio de GitHub.

Agradecemos especialmente a Matthew Mcclain, Beni Hegedus, Kamran Khan, Shruti Koparkar y Qing Lan por revisar y proporcionar aportes valiosos.

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

Share:

Was this article helpful?

93 out of 132 found this helpful

Discover more

Inteligencia Artificial

Conoce a EAGLE Un nuevo método de aprendizaje automático para decodificación rápida de LLM basado en compresión.

Los Grandes Modelos de Lenguaje (LLMs, por sus siglas en inglés) como ChatGPT han revolucionado el procesamiento del ...

Inteligencia Artificial

ChatGPT obtiene una puntuación en el 1% superior en la prueba de creatividad humana

La inteligencia artificial (IA) ha alcanzado nuevas alturas, según una investigación realizada por la Universidad de ...

Inteligencia Artificial

Los investigadores de la Universidad de Pennsylvania presentaron un enfoque alternativo de IA para diseñar y programar computadoras de depósito basadas en RNN.

El cerebro humano es uno de los sistemas más complejos que la naturaleza ha creado. Los neuronas interactúan entre sí...

Inteligencia Artificial

El salto de KPMG hacia el futuro de la IA generativa

En un giro notable de los acontecimientos, el mundo de la consultoría y las finanzas está experimentando un viaje tra...

Ciencia de Datos

Extensiones de AI para Chrome para la hoja de trucos de científicos de datos

La última hoja de trucos de Zepes te presenta una impresionante variedad de herramientas y recursos avanzados diseñad...