Maximice el rendimiento de difusión estable y reduzca los costos de inferencia con AWS Inferentia2 | Servicios web de Amazon

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

En esta publicación, mostramos cómo puede ejecutar modelos de difusión estable y lograr un alto rendimiento al menor costo en Nube informática elástica de Amazon (Amazon EC2) usando Instancias de Amazon EC2 Inf2 creado por Inferencia de AWS2. Observamos la arquitectura de un modelo de difusión estable y recorremos los pasos para compilar un modelo de difusión estable usando AWS neurona y desplegarlo en una instancia Inf2. También analizamos las optimizaciones que Neuron SDK realiza automáticamente para mejorar el rendimiento. Puede ejecutar las versiones Stable Diffusion 2.1 y 1.5 en AWS Inferentia2 de manera rentable. Por último, mostramos cómo puede implementar un modelo de difusión estable en una instancia de Inf2 con Amazon SageMaker.

El tamaño del modelo Stable Diffusion 2.1 en coma flotante 32 (FP32) es de 5 GB y 2.5 GB en bfoat16 (BF16). Una sola instancia inf2.xlarge tiene un acelerador AWS Inferentia2 con 32 GB de memoria HBM. El modelo Stable Diffusion 2.1 puede caber en una sola instancia inf2.xlarge. Stable Diffusion es un modelo de texto a imagen que puede usar para crear imágenes de diferentes estilos y contenido simplemente proporcionando un mensaje de texto como entrada. Para obtener más información sobre la arquitectura del modelo Stable Diffusion, consulte Cree imágenes de alta calidad con modelos Stable Diffusion e impleméntelas de manera rentable con Amazon SageMaker.

Cómo Neuron SDK optimiza el rendimiento de Stable Diffusion

Antes de que podamos implementar el modelo Stable Diffusion 2.1 en las instancias de AWS Inferentia2, debemos compilar los componentes del modelo mediante el SDK de Neuron. Neuron SDK, que incluye un compilador de aprendizaje profundo, tiempo de ejecución y herramientas, compila y optimiza automáticamente los modelos de aprendizaje profundo para que puedan ejecutarse de manera eficiente en instancias Inf2 y extraer el rendimiento completo del acelerador AWS Inferentia2. Tenemos ejemplos disponibles para el modelo Stable Diffusion 2.1 en la Repositorio GitHub. Este cuaderno presenta un ejemplo completo de cómo compilar un modelo de difusión estable, guardar los modelos de Neuron compilados y cargarlos en el tiempo de ejecución para la inferencia.

Utilizamos StableDiffusionPipeline de la cara que abraza diffusers biblioteca para cargar y compilar el modelo. Luego compilamos todos los componentes del modelo para Neuron usando torch_neuronx.trace() y guarde el modelo optimizado como TorchScript. Los procesos de compilación pueden consumir bastante memoria, lo que requiere una cantidad significativa de RAM. Para evitar esto, antes de rastrear cada modelo, creamos un deepcopy de la parte de la tubería que se está rastreando. Después de esto, eliminamos el objeto de tubería de la memoria usando del pipe. Esta técnica es particularmente útil cuando se compila en instancias con poca RAM.

Además, también realizamos optimizaciones a los modelos de difusión estable. UNet tiene el aspecto computacionalmente más intensivo 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 el UNet para un lote (usando tensores de entrada con un lote), luego usamos el torch_neuronx.DataParallel API para cargar este modelo de lote único en cada núcleo. La salida de esta API es un módulo de dos lotes integrado: podemos pasar a UNet las entradas de dos lotes y se devuelve una salida de dos lotes, pero internamente, los dos modelos de un solo lote se ejecutan en los dos núcleos de Neuron. . Esta estrategia optimiza la utilización de recursos y reduce la latencia.

Compile e implemente un modelo de difusión estable en una instancia Inf2 EC2

Para compilar e implementar el modelo Stable Diffusion en una instancia Inf2 EC2, inicie sesión en el Consola de administración de AWS y cree una instancia inf2.8xlarge. Tenga en cuenta que se requiere una instancia inf2.8xlarge solo para la compilación del modelo porque la compilación requiere una memoria de host superior. El modelo Stable Diffusion se puede alojar en una instancia inf2.xlarge. Puede encontrar la última AMI con las bibliotecas de Neuron usando lo siguiente Interfaz de línea de comandos de AWS (AWS CLI) comando:

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 instancia EC2 utilizando Deep Learning AMI Neuron PyTorch 1.13 (Ubuntu 20.04). A continuación, puede crear un entorno de laboratorio de JupyterLab conectándose a la instancia y ejecutando los siguientes pasos:

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

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

Veamos los pasos de compilación para uno de los bloques codificadores de texto. Otros bloques que forman parte de la tubería Stable Diffusion se pueden compilar de manera similar.

El primer paso es cargar el modelo previamente entrenado de Hugging Face. El StableDiffusionPipeline.from_pretrained método carga el modelo pre-entrenado en nuestro objeto de canalización, pipe. Entonces creamos un deepcopy del codificador de texto de nuestra canalización, clonándolo efectivamente. El del pipe Luego, el comando se usa para eliminar el objeto de canalización original, liberando la memoria que consumió. Aquí, estamos cuantificando 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 consiste en envolver nuestro codificador de texto con el NeuronTextEncoder envoltura. La salida de un módulo codificador de texto compilado será de dict. Lo convertimos en un list escriba usando este contenedor:

text_encoder = NeuronTextEncoder(text_encoder)

Inicializamos el tensor PyTorch emb con algunos valores. El emb tensor se utiliza como entrada de ejemplo para el torch_neuronx.trace función. Esta función rastrea nuestro codificador de texto y lo compila en un formato optimizado para Neuron. La ruta del directorio para el modelo compilado se construye uniendo 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 usando torch.jit.save. Se almacena con el nombre de archivo model.pt en el text_encoder directorio 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)

La cuaderno incluye pasos similares para compilar otros componentes del modelo: UNet, decodificador VAE y VAE post_quant_conv. Una vez que haya compilado todos los modelos, puede cargar y ejecutar el modelo siguiendo estos pasos:

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

A continuación, puede ejecutar la canalización proporcionando texto de entrada como indicaciones. Las siguientes son algunas imágenes generadas por el modelo para las indicaciones:

• Retrato de renaud sechan, pluma y tinta, dibujos lineales intrincados, por craig mullins, ruan jia, kentaro miura, greg rutkowski, loundraw

• Retrato de un viejo minero del carbón en el siglo XIX, hermosa pintura, con pintura facial muy detallada de greg rutkowski

• Un castillo en medio de un bosque

Aloje Stable Diffusion 2.1 en AWS Inferentia2 y SageMaker

El alojamiento de modelos Stable Diffusion con SageMaker también requiere compilación con Neuron SDK. Puede completar la compilación antes de tiempo o durante el tiempo de ejecución mediante contenedores de inferencia de modelos grandes (LMI). La compilación anticipada permite tiempos de carga de modelos más rápidos y es la opción preferida.

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

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

Observamos 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 Servicio de almacenamiento simple de Amazon (Amazon S3) cubeta. Puede utilizar este modelo precompilado para sus implementaciones.

Configure el modelo con un script proporcionado

En esta sección, mostramos cómo configurar el contenedor LMI para alojar los modelos de difusión estable. El portátil SD2.1 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 alojar el modelo. La estructura de directorios necesaria es la siguiente:

<config-root-directory> / 
    ├── serving.properties
    │   
    └── model.py [OPTIONAL]

A continuación, creamos el sirviendo.propiedades archivo 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:

• opción.model_id – Los contenedores LMI usan s5cmd para cargar el modelo desde la ubicación S3 y, por lo tanto, debemos especificar la ubicación donde se encuentran nuestros pesos compilados.
• opción.punto de entrada – Para utilizar los controladores integrados, especificamos la clase transformers-neuronx. Si tiene un script de inferencia personalizado, debe proporcionarlo en su lugar.
• opción.dtype – Esto especifica cargar los pesos en un tamaño específico. Para esta publicación, usamos BF16, que reduce aún más nuestros requisitos de memoria frente a FP32 y reduce nuestra latencia debido a eso.
• opción.tensor_parallel_degree – Este parámetro especifica el número de aceleradores que usamos para este modelo. El acelerador de chips AWS Inferentia2 tiene dos núcleos Neuron, por lo que especificar un valor de 2 significa que usamos un acelerador (dos núcleos). Esto significa que ahora podemos crear varios trabajadores para aumentar el rendimiento del punto final.
• opción.motor – Esto está configurado en Python para indicar que no usaremos otros compiladores como DeepSpeed ​​o Faster Transformer para este alojamiento.

Trae tu propio guión

Si desea traer su propia secuencia de comandos de inferencia personalizada, debe eliminar el option.entryPoint Desde serving.properties. El contenedor LMI en ese caso buscará un model.py archivo en la misma ubicación que el serving.properties y usar eso para ejecutar la inferencia.

Cree su propio script de inferencia (model.py)

Crear su propia secuencia de comandos de inferencia es relativamente sencillo utilizando el contenedor LMI. El contenedor requiere su model.py archivo para tener una implementación del siguiente método:

def handle(inputs: Input) which returns an object of type Outputs

Examinemos algunas de las áreas críticas de la cuaderno adjunto, que demuestra la función traer su propia secuencia de comandos.

Vuelva a colocar la cross_attention módulo con la versión optimizada:

# Replace original cross-attention module with custom cross-attention module for better performance
    CrossAttention.get_attention_scores = get_attention_scores
Load the compiled weights for the following
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 del archivo de pesos compilados que usamos al crear las compilaciones. Siéntase libre de cambiar los nombres de los archivos, pero asegúrese de que los nombres de sus archivos de pesos coincidan con lo que especifica aquí.

Luego, debemos cargarlos con Neuron SDK y configurarlos en los pesos reales del modelo. Al cargar los pesos optimizados de UNet, tenga en cuenta que también estamos especificando la cantidad de núcleos Neuron en los que necesitamos cargarlos. Aquí, cargamos a un solo acelerador con dos núcleos:

# Load the compiled UNet onto two neuron cores.
    pipe.unet = NeuronUNet(UNetWrap(pipe.unet))
    logging.info(f"Loading model: unet:created")
    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)
   
 
    # Load other compiled models onto a single neuron core.
 
    # - load encoders
    pipe.text_encoder = NeuronTextEncoder(pipe.text_encoder)
    clip_compiled = torch.jit.load(text_encoder_filename)
    pipe.text_encoder.neuron_text_encoder = clip_compiled
    #- load decoders
    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 aviso invoca el objeto de tubería para generar una imagen.

Crear el punto final de SageMaker

Usamos las API de Boto3 para crear un punto final de SageMaker. Complete los siguientes pasos:

1. Cree el tarball con solo la porción y el opcional model.py y cárguelo en Amazon S3.
2. Cree el modelo utilizando el contenedor de imágenes y el tarball del modelo cargado anteriormente.
3. Cree la configuración del punto final utilizando los siguientes parámetros clave:
   1. Usa una ml.inf2.xlarge ejemplo.
   2. Set ContainerStartupHealthCheckTimeoutInSeconds a 240 para garantizar que la verificación de estado comience después de implementar el modelo.
   3. Set VolumeInGB a un valor mayor para que pueda usarse 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 cargarlo en Amazon S3, debemos crear un modelo de SageMaker. Usamos 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 por defecto. Ver el siguiente código:

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

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

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 final de SageMaker

En esta demostración, usamos una instancia ml.inf2.xlarge. Necesitamos establecer el VolumeSizeInGB parámetros para proporcionar el espacio en disco necesario para cargar el modelo y los pesos. Este parámetro es aplicable a las instancias que soportan el Tienda de bloques elásticos de Amazon (Amazon EBS) archivo adjunto de volumen. Podemos dejar el tiempo de espera de descarga del modelo y la verificación del estado de inicio del contenedor en un valor más alto, lo que le dará el tiempo adecuado para que el contenedor extraiga los pesos de Amazon S3 y se cargue en los aceleradores de AWS Inferentia2. Para obtener más detalles, consulte CreateEndpointConfigCreateEndpointConfig.

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 final de SageMaker:

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

Invocar el punto final del modelo

Este es un modelo generativo, por lo que pasamos el aviso que el modelo usa para generar la imagen. El payload es del tipo JSON:

response_model = boto3_sm_run_client.invoke_endpoint( EndpointName=endpoint_name,
    Body=json.dumps(
        {
            "prompt": "Mountain Landscape", 
            "parameters": {} # 
        }
    ), 
    ContentType="application/json",
)

Evaluación comparativa del modelo de difusión estable en Inf2

Realizamos algunas pruebas para comparar el modelo Stable Diffusion con el tipo de datos BF 16 en Inf2, y pudimos derivar números de latencia que rivalizan o superan algunos de los otros aceleradores para Stable Diffusion. Esto, junto con el menor costo de los chips AWS Inferentia2, hace que esta sea 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 Instancias Inf2 de Amazon EC2.

Conclusión

En esta publicación, profundizamos en la compilación, optimización e 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 una excelente relación precio-rendimiento para Stable Diffusion 1.5. Para obtener más información sobre por qué las instancias Inf2 son excelentes para la IA generativa y los modelos de lenguaje extenso, consulte Las instancias Inf2 de Amazon EC2 para la inferencia de IA generativa de bajo costo y alto rendimiento ya están disponibles de forma general. Para obtener detalles sobre el rendimiento, consulte Rendimiento Inf2. Vea ejemplos adicionales en el Repositorio GitHub.

Un agradecimiento especial a Matthew McClain, Beni Hegedus, Kamran Khan, Shruti Koparkar y Qing Lan por revisar y brindar valiosos aportes.

Acerca de los autores

Vivek Gangasani es arquitecto sénior de soluciones de aprendizaje automático en Amazon Web Services. Trabaja con nuevas empresas de aprendizaje automático para crear e implementar aplicaciones de IA/ML en AWS. Actualmente se enfoca en brindar soluciones para MLOps, inferencia de ML y ML de código bajo. Ha trabajado en proyectos en diferentes dominios, incluido el procesamiento del lenguaje natural y la visión artificial.

kc tung es arquitecto sénior de soluciones en AWS Annapurna Labs. Se especializa en el entrenamiento y la implementación a gran escala de modelos de aprendizaje profundo en la nube. Tiene un doctorado. en biofísica molecular del Centro Médico Southwestern de la Universidad de Texas en Dallas. Ha hablado en AWS Summits y AWS Reinvent. Hoy ayuda a los clientes a entrenar e implementar grandes modelos de PyTorch y TensorFlow en la nube de AWS. Es autor de dos libros: Aprende TensorFlow Enterprise y Referencia de bolsillo de TensorFlow 2.

Rupinder Grewal es un Arquitecto de Soluciones Sr. Ai/ML Especialista con AWS. Actualmente se enfoca en servir modelos y MLOps en SageMaker. Antes de este cargo, trabajó como ingeniero de aprendizaje automático creando y alojando modelos. Fuera del trabajo, le gusta jugar tenis y andar en bicicleta por senderos de montaña.

• Distribución de relaciones públicas y contenido potenciado por SEO. Consiga amplificado hoy.
• PlatoData.Network Vertical Generativo Ai. Empodérate. Accede Aquí.
• PlatoAiStream. Inteligencia Web3. Conocimiento amplificado. Accede Aquí.
• PlatoESG. Automoción / vehículos eléctricos, Carbón, tecnología limpia, Energía, Ambiente, Solar, Gestión de residuos. Accede Aquí.
• Desplazamientos de bloque. Modernización de la propiedad de compensaciones ambientales. Accede Aquí.
• Fuente: https://aws.amazon.com/blogs/machine-learning/maximize-stable-diffusion-performance-and-lower-inference-costs-with-aws-inferentia2/