Pixart-sigma es un transformador de difusión modelo que es capaz de generación de imágenes en resolución 4K. Este modelo muestra mejoras significativas sobre los modelos Pixart de generación anteriores como Pixart-Alpha y otros modelos de difusión a través de mejoras de datos y mejoras arquitectónicas. AWS Entrenium y AWS Inferentia son chips de IA especialmente diseñados para acelerar las cargas de trabajo de aprendizaje automático (ML), lo que las hace ideales para el despliegue rentable de modelos generativos grandes. Al usar estos chips AI, puede lograr un rendimiento y eficiencia óptimos al ejecutar la inferencia con modelos de transformadores de difusión como Pixart-Sigma.
Esta publicación es la primera de una serie en la que ejecutaremos múltiples transformadores de difusión en casos de entrenamiento e inferentios. En esta publicación, mostramos cómo puede implementar Pixart-Sigma en capacitación y instancias propuladas por inferentes.
Descripción general de la solución
Los pasos que se describen a continuación se utilizarán para implementar el modelo Pixart-Sigma en AWS Entrenium y ejecutar inferencia para generar imágenes de alta calidad.
- Paso 1-Requisitos previos y configuración
- Paso 2: descargue y compile el modelo Pixart-Sigma para AWS Entrenium
- Paso 3: implementa el modelo en AWS Entrenium para generar imágenes
Paso 1 – Requisitos previos y configuración
Para comenzar, deberá configurar un entorno de desarrollo en un host TRN1, TRN2 o INF2. Complete los siguientes pasos:
- Lanzar un
trn1.32xlargeotrn2.48xlargeinstancia con una neurona dlami. Para obtener instrucciones sobre cómo comenzar, consulte Comience con Neuron en Ubuntu 22 con Neuron Multi-Framework Dlami. - Lanza un cuaderno Jupyter sever. Para obtener instrucciones para configurar un servidor Jupyter, consulte lo siguiente guía del usuario.
- Clonar el AWS-Neuron-Samples Repositorio de GitHub:
- Navegar al hf_preetrainse_pixart_sigma_1k_latency_optimized.ipynb computadora portátil:
El script de ejemplo proporcionado está diseñado para ejecutarse en una instancia de TRN2, pero puede adaptarlo para instancias TRN1 o INF2 con modificaciones mínimas. Específicamente, dentro del cuaderno y en cada uno de los archivos de componentes bajo el neuron_pixart_sigma Directorio, encontrará cambios comentados para acomodar configuraciones TRN1 o INF2.
Paso 2: descargue y compile el modelo Pixart-Sigma para AWS Entrenium
Esta sección proporciona una guía paso a paso para compilar Pixart-Sigma para AWS Entrenium.
Descarga el modelo
Encontrarás una función de ayuda en caché-hf-model.py En el repositorio de GitHub mencionado anteriormente que muestra cómo descargar el modelo Pixart-Sigma de Hugging Face. Si está utilizando Pixart-Sigma en su propia carga de trabajo y opta por no usar el script incluido en esta publicación, puede usar el Huggingface-cli para descargar el modelo en su lugar.
La implementación Neuron Pixart-Sigma contiene algunos scripts y clases. Los diversos archivos y scrips se descomponen de la siguiente manera:
├── compile_latency_optimized.sh # Full Model Compilation script for Latency Optimized
├── compile_throughput_optimized.sh # Full Model Compilation script for Throughput Optimized
├── hf_pretrained_pixart_sigma_1k_latency_optimized.ipynb # Notebook to run Latency Optimized Pixart-Sigma
├── hf_pretrained_pixart_sigma_1k_throughput_optimized.ipynb # Notebook to run Throughput Optimized Pixart-Sigma
├── neuron_pixart_sigma
│ ├── cache_hf_model.py # Model downloading Script
│ ├── compile_decoder.py # Text Encoder Compilation Script and Wrapper Class
│ ├── compile_text_encoder.py # Text Encoder Compilation Script and Wrapper Class
│ ├── compile_transformer_latency_optimized.py # Latency Optimized Transformer Compilation Script and Wrapper Class
│ ├── compile_transformer_throughput_optimized.py # Throughput Optimized Transformer Compilation Script and Wrapper Class
│ ├── neuron_commons.py # Base Classes and Attention Implementation
│ └── neuron_parallel_utils.py # Sharded Attention Implementation
└── requirements.txtEste cuaderno lo ayudará a descargar el modelo, compilar los modelos de componentes individuales e invocar la tubería de generación para generar una imagen. Aunque los cuadernos se pueden ejecutar como una muestra independiente, las siguientes secciones de esta publicación se realizarán a través de los detalles clave de implementación dentro de los archivos y scripts de componentes para admitir la ejecución de Pixart-Sigma en Neuron.
Para cada componente de Pixart (T5, Transformer y VAE), el ejemplo utiliza clases de envoltura específicas de neuronas. Estas clases de envoltura tienen dos propósitos. El primer propósito es que nos permite rastrear los modelos para la compilación:
class InferenceTextEncoderWrapper(nn.Module):
def __init__(self, dtype, t: T5EncoderModel, seqlen: int):
super().__init__()
self.dtype = dtype
self.device = t.device
self.t = t
def forward(self, text_input_ids, attention_mask=None):
return [self.t(text_input_ids, attention_mask)['last_hidden_state'].to(self.dtype)]
Por favor consulte el neuron_commons.py Archivo para todos los módulos y clases de envoltura.
La segunda razón para usar clases de envoltura es modificar la implementación de atención para que se ejecute en Neuron. Debido a que los modelos de difusión como Pixart generalmente están unidos al cómputo, puede mejorar el rendimiento fragmentando la capa de atención en múltiples dispositivos. Para hacer esto, reemplaza las capas lineales con neuronx distribuido RowParallellineal y Columnparallellineal Capas:
def shard_t5_self_attention(tp_degree: int, selfAttention: T5Attention):
orig_inner_dim = selfAttention.q.out_features
dim_head = orig_inner_dim // selfAttention.n_heads
original_nheads = selfAttention.n_heads
selfAttention.n_heads = selfAttention.n_heads // tp_degree
selfAttention.inner_dim = dim_head * selfAttention.n_heads
orig_q = selfAttention.q
selfAttention.q = ColumnParallelLinear(
selfAttention.q.in_features,
selfAttention.q.out_features,
bias=False,
gather_output=False)
selfAttention.q.weight.data = get_sharded_data(orig_q.weight.data, 0)
del(orig_q)
orig_k = selfAttention.k
selfAttention.k = ColumnParallelLinear(
selfAttention.k.in_features,
selfAttention.k.out_features,
bias=(selfAttention.k.bias is not None),
gather_output=False)
selfAttention.k.weight.data = get_sharded_data(orig_k.weight.data, 0)
del(orig_k)
orig_v = selfAttention.v
selfAttention.v = ColumnParallelLinear(
selfAttention.v.in_features,
selfAttention.v.out_features,
bias=(selfAttention.v.bias is not None),
gather_output=False)
selfAttention.v.weight.data = get_sharded_data(orig_v.weight.data, 0)
del(orig_v)
orig_out = selfAttention.o
selfAttention.o = RowParallelLinear(
selfAttention.o.in_features,
selfAttention.o.out_features,
bias=(selfAttention.o.bias is not None),
input_is_parallel=True)
selfAttention.o.weight.data = get_sharded_data(orig_out.weight.data, 1)
del(orig_out)
return selfAttention
Por favor consulte el neuron_parallel_utils.py Presente para obtener más detalles sobre la atención paralela.
Compilar submodelos individuales
El modelo Pixart-Sigma está compuesto por tres componentes. Cada componente se compila para que toda la tubería de generación pueda ejecutarse en Neuron:
- Codificador de texto -Un codificador de 4 billones de parámetros, que traduce un aviso legible por humanos en una incrustación. En el codificador de texto, las capas de atención están fragmentadas, junto con las capas de alimentación hacia adelante, con paralelismo tensor.
- Modelo de transformador de Denoising -Un transformador de 700 millones de parámetros, que iterativamente denoja un latente (una representación numérica de una imagen comprimida). En el transformador, las capas de atención están fiscasas, junto con las capas de alimentación hacia adelante, con paralelismo tensor.
- Descifrador -Un decodificador de VAE que convierte nuestro latente generado por Denoiser a una imagen de salida. Para el decodificador, el modelo se implementa con el paralelismo de datos.
Ahora que la definición del modelo está lista, debe rastrear un modelo para ejecutarlo en entrenamiento o inferencia. Puedes ver cómo usar el trace() Función para compilar el modelo de componente del decodificador para Pixart en el siguiente bloque de código:
compiled_decoder = torch_neuronx.trace(
decoder,
sample_inputs,
compiler_workdir=f"{compiler_workdir}/decoder",
compiler_args=compiler_flags,
inline_weights_to_neff=False
)
Por favor consulte el compile_decoder.py Presente más sobre cómo instanciar y compilar el decodificador.
Ejecutar modelos con paralelismo tensoruna técnica utilizada para dividir un tensor en trozos en múltiples neuroncores, debe rastrear con un tp_degree. Este tp_degree Especifica el número de neuroncoros para fastar el modelo. Luego usa el parallel_model_trace API para compilar los modelos de componentes codificadores y transformadores para Pixart:
compiled_text_encoder = neuronx_distributed.trace.parallel_model_trace(
get_text_encoder_f,
sample_inputs,
compiler_workdir=f"{compiler_workdir}/text_encoder",
compiler_args=compiler_flags,
tp_degree=tp_degree,
)
Por favor consulte el compile_text_encoder.py Presente para obtener más detalles sobre el rastreo del codificador con paralelismo tensor.
Por último, trazas el modelo de transformador con paralelismo tensor:
compiled_transformer = neuronx_distributed.trace.parallel_model_trace(
get_transformer_model_f,
sample_inputs,
compiler_workdir=f"{compiler_workdir}/transformer",
compiler_args=compiler_flags,
tp_degree=tp_degree,
inline_weights_to_neff=False,
)
Por favor consulte el compile_transformer_latency_optimized.py Presente para obtener más detalles sobre el rastreo del transformador con paralelismo tensor.
Usarás el compile_latency_optimized.sh Script para compilar los tres modelos como se describe en esta publicación, por lo que estas funciones se ejecutarán automáticamente cuando se ejecute a través del cuaderno.
Paso 3: implementa el modelo en AWS Entrenium para generar imágenes
Esta sección nos guiará a través de los pasos para ejecutar inferencia en Pixart-Sigma en AWS Entrenium.
Crear un objeto de tubería de difusores
La biblioteca de difusores de la cara abrazada es una biblioteca para modelos de difusión previamente capacitados e incluye tuberías específicas del modelo que agrupan los componentes (modelos, programadores y procesadores entrenados independientemente) necesarios para ejecutar un modelo de difusión. El PixArtSigmaPipeline es específico del modelo Pixartsigma y se instancia de la siguiente manera:
pipe: PixArtSigmaPipeline = PixArtSigmaPipeline.from_pretrained(
"PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
torch_dtype=torch.bfloat16,
local_files_only=True,
cache_dir="pixart_sigma_hf_cache_dir_1024")
Por favor consulte el hf_preetrainse_pixart_sigma_1k_latency_optimized.ipynb cuaderno para obtener detalles sobre la ejecución de la tubería.
Cargar modelos de componentes compilados en la tubería de generación
Después de que se haya compilado cada modelo de componente, cargue en la tubería de generación general para la generación de imágenes. El modelo VAE está cargado de paralelismo de datos, lo que nos permite paralelizar la generación de imágenes para el tamaño de lotes o múltiples imágenes por aviso. Para más detalles, consulte el hf_preetrainse_pixart_sigma_1k_latency_optimized.ipynb computadora portátil.
vae_decoder_wrapper.model = torch_neuronx.DataParallel(
torch.jit.load(decoder_model_path), [0, 1, 2, 3], False
)
text_encoder_wrapper.t = neuronx_distributed.trace.parallel_model_load(
text_encoder_model_path
)
Finalmente, los modelos cargados se agregan a la tubería de generación:
pipe.text_encoder = text_encoder_wrapper
pipe.transformer = transformer_wrapper
pipe.vae.decoder = vae_decoder_wrapper
pipe.vae.post_quant_conv = vae_post_quant_conv_wrapper
Componer un aviso
Ahora que el modelo está listo, puede escribir un mensaje para transmitir qué tipo de imagen desea generar. Al crear un aviso, siempre debe ser lo más específico posible. Puede usar un mensaje positivo para transmitir lo que se desea en su nueva imagen, incluido un sujeto, acción, estilo y ubicación, y puede usar un indicador negativo para indicar características que deben eliminarse.
Por ejemplo, puede usar las siguientes indicaciones positivas y negativas para generar una foto de un astronauta que monta un caballo en Marte sin montañas:
# Subject: astronaut
# Action: riding a horse
# Location: Mars
# Style: photo
prompt = "a photo of an astronaut riding a horse on mars"
negative_prompt = "mountains"
Siéntase libre de editar el aviso en su cuaderno usando ingeniería rápida para generar una imagen de su elección.
Generar una imagen
Para generar una imagen, pasa la solicitud a la tubería del modelo Pixart y luego guarda la imagen generada para una referencia posterior:
# pipe: variable holding the Pixart generation pipeline with each of
# the compiled component models
images = pipe(
prompt=prompt,
negative_prompt=negative_prompt,
num_images_per_prompt=1,
height=1024, # number of pixels
width=1024, # number of pixels
num_inference_steps=25 # Number of passes through the denoising model
).images
for idx, img in enumerate(images):
img.save(f"image_{idx}.png")
Limpieza
Para evitar incurrir en costos adicionales, Detén su instancia de EC2 Usando o Consola de gestión de AWS o Interfaz de línea de comandos de AWS (AWS CLI).
Conclusión
En esta publicación, caminamos por cómo desplegar Pixart-Sigma, un transformador de difusión de última generación, en instancias de entrenamiento. Esta publicación es la primera de una serie centrada en ejecutar transformadores de difusión para diferentes tareas de generación en Neuron. Para obtener más información sobre la ejecución de modelos de transformadores de difusión con neurona, consulte Transformadores de difusión.
Sobre los autores
Achintya pinninti es un arquitecto de soluciones en Amazon Web Services. Apoya a los clientes del sector público, lo que les permite lograr sus objetivos utilizando la nube. Se especializa en la creación de datos y soluciones de aprendizaje automático para resolver problemas complejos.
Miriam Lebowitz es un arquitecto de soluciones centrado en empoderar a las nuevas empresas en la etapa inicial en AWS. Ella aprovecha su experiencia con AI/ML para guiar a las empresas para seleccionar e implementar las tecnologías adecuadas para sus objetivos comerciales, estableciéndolos para un crecimiento e innovación escalables en el mundo de inicio competitivo.
Sadaf rasool es un arquitecto de soluciones en Annapurna Labs en AWS. SADAF colabora con los clientes para diseñar soluciones de aprendizaje automático que aborden sus desafíos comerciales críticos. Ayuda a los clientes a capacitar e implementar modelos de aprendizaje automático que aprovechen los chips de AWS Trainium o AWS Inferentia para acelerar su viaje de innovación.
John Gray es un arquitecto de soluciones en Annapurna Labs, AWS, con sede en Seattle. En este rol, John trabaja con los clientes en sus casos de uso de AI y aprendizaje automático, Architects Solutions para resolver de manera rentable sus problemas comerciales y les ayuda a construir un prototipo escalable utilizando chips AWS AI.