Cómo entrené 10 TB para difusión estable en SageMaker

¡Con detalles sobre mi próximo libro sobre capacitación distribuida con AWS!

Tratar con muchos datos es difícil, no hay duda al respecto. Sin embargo, con algunas elecciones inteligentes y diseños simples, usted también puede ampliar sus proyectos para replicar y producir resultados de última generación.

En esta publicación, lo guiaré a través de un caso de estudio en el que utilicé AWS, especialmente SageMaker, S3 y FSx for Lustre, para descargar, procesar y entrenar un modelo de difusión estable. Usé SageMaker Distributed Data Parallel para optimizar el descenso de gradiente a escala. También utilicé el paralelismo de trabajos en el backend modular de SageMaker para descargar y procesar todos los archivos.

Mi código de capacitación se basa en una muestra de Corey Barrett de AWS MLSL, disponible aquí. Intentaré actualizar ese repositorio con todo el código fuente de este proyecto para fin de año.

Si está ansioso por obtener más contenido sobre la ejecución de trabajos a gran escala en SageMaker y el alcance completo de la capacitación previa de modelos de visión y lenguaje grandes en AWS, consulte mi libro que se publicará en abril de 2023.

En mi libro, dedico 15 capítulos a guiarlo a través de todo el ciclo de vida, desde encontrar el caso de uso correcto hasta preparar su entorno, capacitar, solucionar problemas, evaluar, poner en funcionamiento y escalar.

Mi libro de entrenamiento distribuido es ideal para principiantes ; es realmente para desarrolladores principiantes e intermedios de Python que desean eventualmente adquirir las habilidades para entrenar sus propios modelos de visión y lenguaje amplios.

¡Vamos a sumergirnos!

1/ Descargar archivos de parquet Laion-5B con trabajos de SageMaker

El conjunto de datos central utilizado para entrenar Stable Diffusion es Laion-5B. Este es un conjunto de datos de código abierto que proporciona miles de millones de pares de imagen/texto de Internet. Romain Beaumont, colaborador del conjunto de datos Laion-5B, abrió su herramienta para descargar y procesar este conjunto de datos, y lo usé como base. Su código fuente está disponible aquí. ¡Gracias Romaín!

El conjunto de datos central consta de dos partes: más de 100 archivos de parquet que apuntan a las imágenes originales y sus leyendas, y las propias imágenes descargadas.

Primero, descargué los archivos de parquet con una simple función de Python. Afortunadamente, el kit de herramientas de Romain le permite descargar directamente a su depósito S3, lo que le permite evitar cualquier restricción de almacenamiento local. Probé esto varias veces en mi instancia de SageMaker Studio, saltando a una máquina más grande cuando necesitaba trabajar con un archivo más grande.

Después de eso, lo envolví en un único script que podía ejecutar en mi trabajo de SageMaker.

import argparse
import os

def parse_args():
    
    parser = argparse.ArgumentParser()    
    
    # points to your session bucket                 
    parser.add_argument("--bucket", type=str, default=os.environ["SM_HP_BUCKET"])
    
    args = parser.parse_args()
    
    return args

def download_parquet(bucket):
    cmd = '''for i in {00000..00127}; 
            do wget https://huggingface.co/datasets/laion/laion2B-en/resolve/main/part-$i-5114fd87-297e-42b0-9d11-50f1df323dfa-c000.snappy.parquet -O - | aws s3 cp - s3://{}/metadata/laion2B-en-joined/part-$i-4cfd6e30-f032-46ee-9105-8696034a8373-c000.snappy.parquet;  
        done'''.format(bucket)

    # runs the command
    os.system(cmd)

if __name__ == "__main__":
    
    args = parse_args()
        
    download_parquet(args.bucket)

2/ Descargar pares de imagen y texto de Laion-5B

Después de esto, llegó el momento de descargar las imágenes en sí. Hice una lista de todos los archivos de parquet, la revisé y envié cada archivo de parquet a su propio trabajo de SageMaker.

Mi secuencia de comandos para descargar los pares de imagen y texto se veía así:

from img2dataset import download
import shutil
import os
import multiprocessing
import threading
import argparse

def parse_args():
    
    parser = argparse.ArgumentParser()
        
    parser.add_argument("--cores", type=int, default=multiprocessing.cpu_count())

    parser.add_argument("--threads", type=int, default=threading.active_count())
    
    parser.add_argument("--parquet", type=str, default=os.environ["SM_CHANNEL_PARQUET"])
    
    parser.add_argument("--file_name", type=str, default=os.environ["SM_HP_FILE_NAME"])
            
    parser.add_argument("--bucket", type=str, default=os.environ["SM_MODULE_DIR"].split('/')[2])
        
    args = parser.parse_args()
    
    return args

def prep_system():
    
    args = parse_args()
    
    # send joint path and file name
    url_list = "{}/{}".format(args.parquet, args.file_name)
    
    part_number = args.file_name.split('-')[1]

    # point to output path in S3
    s3_output = "s3://{}/data/part-{}/".format(args.bucket, part_number)
    
    return args, url_list, s3_output

    
if __name__ == "__main__":
    
    args, url_list, s3_output = prep_system()
    
    download(
        processes_count=args.cores,
        thread_count=args.threads,
        # takes a single parquet file
        url_list=url_list,
        image_size=256,
        # copies to S3 directly, bypassing local disk
        output_folder=s3_output,
        # each image / caption pair is a tarball
        output_format="webdataset",
        input_format="parquet",
        url_col="URL",
        caption_col="TEXT",
        enable_wandb=False,
        number_sample_per_shard=1000,
        distributor="multiprocessing",
    )

Luego, revisé mi lista de archivos de parquet y envié cada archivo de parquet a su propio trabajo . A esto lo llamamos paralelismo de trabajos , es una manera fácil y efectiva de ejecutar trabajos de cómputo paralelos.

def get_estimator(part_number, p_file, output_dir):
    
    hyperparameters = {"file_name": p_file}

    estimator = PyTorch(entry_point="download.py",
                          base_job_name="laion-part-{}".format(part_number),
                          role=role,
                          source_dir="scripts",
                          # configures the SageMaker training resource, you can increase as you need
                          instance_count=1,
                          instance_type="ml.c5.18xlarge",
                          py_version="py36",
                          framework_version = '1.8',
                          sagemaker_session=sagemaker_session,
                          volume_size = 250,
                          debugger_hook_config=False,
                          hyperparameters=hyperparameters,
                          output_path = output_dir)
    return estimator

for p_file in parquet_list[:18]:
    
    part_number = p_file.split('-')[1]

    output_dir = "s3://{}/data/part-{}/".format(bucket, part_number)

    if is_open(output_dir):

        est = get_estimator(part_number, p_file, output_dir)

        est.fit({"parquet":"s3://{}/meta/{}".format(bucket, p_file)}, wait=False)

Ejecuté esto en 18 trabajos diferentes, todos ejecutándose al mismo tiempo, durante aproximadamente 24 horas.

Los objetos resultantes se veían así:

3/ Crear volumen FSx for Lustre desde la ruta S3

A continuación, dediqué la friolera de veinte minutos a crear un nuevo volumen FSx for Lustre a partir de mis rutas S3. Esto fue sorprendentemente fácil de hacer; ida y vuelta, incluidos algunos de mis propios contratiempos, y el tiempo de creación del volumen fue de unos veinte minutos. Todo sin salir de la consola de AWS.

S3 tiene una característica muy buena para calcular el tamaño total de una ruta. Usé eso para estimar la cantidad total de datos que estaba usando: unos modestos 9,5 TB.

Mi configuración de Lustre se veía así:

¿Por qué aprovisionar Lustre? Porque reduce el tiempo de descarga de datos en su trabajo de capacitación de SageMaker, a esta escala probablemente unos sólidos 90 minutos o más, a menos de 60 segundos.

Pasé algún tiempo configurando mi VPC para habilitar el entrenamiento en SageMaker y escribir los resultados en S3. También amplí un contenedor de aprendizaje profundo de AWS con todos mis paquetes para ahorrar tiempo durante la capacitación; en lugar de instalar pip todo para cada trabajo, lo hago una vez cuando construyo la imagen, luego simplemente descargo esa imagen en mi trabajo en tiempo de ejecución.

4/ Construir un índice de líneas json

Después de esto, pasé algunas horas tratando de disputar la función load_dataset() preconstruida con mi formato webdataset. Todavía sospecho que esto es posible, pero opté por escribir mi propio cargador de datos.

Para ahorrar tiempo y dinero en mi enorme clúster de GPU, construí un índice de líneas json . Se trataba simplemente de un archivo de varios GB con 50 millones de objetos json, cada uno de los cuales apuntaba a un par de imagen/texto válido en FSx para Lustre.

Ejecuté un gran trabajo de SageMaker basado en CPU para enumerar todos los archivos de mis 18 partes en Lustre. Después de esto, cargué, probé y construí una función de cargador de datos muy pequeña en mi computadora portátil Studio.

5/ Ejecutar en 192 GPU con capacitación distribuida de SageMaker

¡Después de eso, era hora de superar una sola época! Aumenté mis límites para las instancias ml.p4d.24xlarge en el entrenamiento de SageMaker en us-east-1 usando la cuota de servicio de AWS a 24. Dado que cada instancia tiene 8 GPU, es un total de 192 GPU.

Hice algunas ejecuciones con solo el 1% de mis datos generales para asegurarme de que el ciclo de entrenamiento funcionara como se esperaba; ¡ La nueva función de piscinas cálidas fue extremadamente útil!

Mi configuración de entrenamiento final se veía así.

import sagemaker
from sagemaker.huggingface import HuggingFace

sess = sagemaker.Session()

role = sagemaker.get_execution_role()

version = 'v1'

image_uri = '<aws account id>.dkr.ecr.us-east-1.amazonaws.com/stable-diffusion:{}'.format(version )

# required in this version of the train script
data_channels['sd_base_model'] = 's3://dist-train/stable-diffusion/conceptual_captions/sd-base-model/'

hyperparameters={'pretrained_model_name_or_path':'/opt/ml/input/data/sd_base_model',
                'train_data_dir':'/opt/ml/input/data/training/laion-fsx',
                'index_name':'data_index.jsonl',
                'caption_column':'caption',
                'image_column':'image',
                'resolution':256,
                'mixed_precision':'fp16',
                # this is per device
                'train_batch_size':22,
                'learning_rate': '1e-10',
                # 'max_train_steps':1000000,
                'num_train_epochs':1,
                'output_dir':'/opt/ml/model/sd-output-final',  
                'n_rows':50000000}

est = HuggingFace(entry_point='finetune.py',
                  source_dir='src',
                  image_uri=image_uri,
                  sagemaker_session=sess,
                  role=role,
                  output_path="s3://laion-5b/output/model/", 
                  instance_type='ml.p4dn.24xlarge',
                  keep_alive_period_in_seconds = 60*60,
                  py_version='py38',
                  base_job_name='fsx-stable-diffusion', 
                  instance_count=24,
                  enable_network_isolation=True,
                  encrypt_inter_container_traffic = True,
                  # all opt/ml paths point to SageMaker training 
                  hyperparameters = hyperparameters,
                  distribution={"smdistributed": { "dataparallel": { "enabled": True } }},
                  max_retry_attempts = 30,
                  max_run = 4 * 60 * 60,
                  debugger_hook_config=False,
                  disable_profiler = True,
                  **kwargs)

est.fit(inputs=data_channels, wait=False)

Lo que más me sorprendió fue que la diferencia entre mis trabajos más grandes y más pequeños, de 7 millones a 50 millones de pares de imagen/texto, fue de solo unos 20 minutos. Ambos trabajos utilizaron 24 instancias, y ambos dedicaron la mayor parte de su tiempo simplemente a inicializar las GPU para el trabajo. Descargar la imagen, configurar MPI, cargar los datos y preparar NCCL llevó la mayor parte del tiempo de ejecución.

¿Mi tiempo para completar 1 época en 50 millones de pares de imagen/texto con 200 GPU? ¡¡15 minutos!!

Así es, 15 minutos. De las 18:30 a las 18:45, la única vez que se utilizaron realmente mis GPU. El trabajo general tardó aproximadamente una hora en completarse, y nuevamente la mayor parte del tiempo se dedicó a inicializar, aprovisionar, cargar y cargar el modelo terminado.

Para mí, esto indica una gran oportunidad en el diseño y optimización de sistemas de GPU. ¿Necesitamos aceptar ese impuesto de 40 minutos solo para entrenar a nuestros modelos? ¡De ninguna manera! Seamos creativos con las soluciones para solucionar esto.

Para ser claros, no entrené este modelo para la convergencia. Simplemente probé el rendimiento en una sola época, que registré a los 15 minutos. Para entrenar durante 1000 épocas, presupuestaría 15 000 minutos, lo que debería tomar un poco más de 10 días.

7/ Implementar en el alojamiento de Amazon SageMaker

Después de esto, implementé el modelo base en el alojamiento de SageMaker. Como es de esperar, un modelo entrenado en una sola época no funcionará muy bien, así que simplemente usé la primera versión de Stable Diffusion en mi propia cuenta de AWS.

¡Y aquí está mi imagen resultante!

¡Y eso es una envoltura! En esta publicación, expliqué cómo trabajé con 10 TB de datos para finalmente entrenar una sola época de Stable Diffusion en el entrenamiento distribuido de SageMaker. ¡Mi ciclo de entrenamiento se completó en 15 minutos!

Espero haber comenzado a despertar su apetito por más contenido sobre el entrenamiento previo de modelos de lenguaje y visión amplia. Estoy escribiendo un libro sobre este tema , que estará disponible en abril de 2023. ¡Podrás unirte a mi Círculo de Autores en enero, asistir a la fiesta de lanzamiento, reunirte con otros expertos en esta área y más!