Como treinei 10 TB para difusão estável no SageMaker

Com detalhes sobre meu próximo livro sobre treinamento distribuído com AWS!

Lidar com muitos dados é difícil, não há dúvida. No entanto, com algumas escolhas inteligentes e designs simples, você também pode ampliar seus projetos para replicar e produzir resultados de última geração.

Nesta postagem, mostrarei um estudo de caso em que usei AWS, principalmente SageMaker, S3 e FSx para Lustre, para baixar, processar e treinar um modelo de difusão estável. Usei o SageMaker Distributed Data Parallel para otimizar a descida do gradiente em escala. Também usei o paralelismo de trabalho no back-end modular do SageMaker para baixar e processar todos os arquivos.

Meu código de treinamento é baseado em uma amostra de Corey Barrett da AWS MLSL, disponível aqui. Vou tentar atualizar esse repositório com todo o código-fonte deste projeto até o final do ano.

Se você deseja obter mais conteúdo sobre a execução de trabalhos de grande escala no SageMaker e todo o escopo do pré-treinamento de modelos de linguagem e visão ampla na AWS, confira meu livro que será lançado em abril de 2023!

Em meu livro, passo 15 capítulos orientando você por todo o ciclo de vida, desde encontrar o caso de uso certo até preparar seu ambiente, treinamento, solução de problemas, avaliação, operacionalização e dimensionamento.

Meu livro de treinamento distribuído é ideal para iniciantes ; é realmente para desenvolvedores Python iniciantes a intermediários que desejam eventualmente adquirir as habilidades para treinar sua própria visão ampla e modelos de linguagem.

Vamos mergulhar!

1/ Baixar arquivos parquet Laion-5B com trabalhos SageMaker

O conjunto de dados principal usado para treinar a difusão estável é o Laion-5B. Este é um conjunto de dados de código aberto que fornece bilhões de pares de imagem/texto da Internet. Romain Beaumont, um colaborador do conjunto de dados Laion-5B, abriu o código de sua ferramenta para baixar e processar esse conjunto de dados, e eu usei isso como minha base. Seu código-fonte está disponível aqui. Obrigado Romain!

O conjunto de dados principal é composto por duas partes: mais de 100 arquivos em parquet que apontam para as imagens originais e suas legendas e as próprias imagens baixadas.

Primeiro, baixei os arquivos parquet com uma função Python simples. Felizmente, o kit de ferramentas do Romain permite que você baixe diretamente para o seu bucket S3, o que permite contornar quaisquer restrições de armazenamento local. Testei isso algumas vezes em minha instância do SageMaker Studio, pulando para uma máquina maior quando precisava trabalhar com um arquivo maior.

Depois disso, envolvi-o em um único script que poderia ser executado em meu trabalho do 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/ Baixe os pares de imagem e texto Laion-5B

Depois disso, era hora de baixar as próprias imagens. Fiz uma lista de todos os arquivos de parquet, examinei-a e enviei cada arquivo de parquet para seu próprio trabalho no SageMaker.

Meu script para baixar os pares de imagem e texto ficou assim:

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",
    )

Em seguida, examinei minha lista de arquivos de parquet e enviei cada arquivo de parquet para seu próprio trabalho . Chamamos isso de paralelismo de trabalho , é uma maneira fácil e eficaz de executar trabalhos de computação 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)

Executei isso em 18 trabalhos diferentes, todos ao mesmo tempo, por cerca de 24 horas.

Os objetos resultantes ficaram assim:

3/ Crie o volume FSx for Luster a partir do caminho S3

Em seguida, gastei impressionantes vinte minutos para criar um novo volume FSx for Lustre a partir de meus caminhos S3. Isso foi surpreendentemente fácil de fazer; ida e volta, incluindo alguns dos meus próprios soluços e o tempo de criação do volume foi de cerca de vinte minutos. Tudo sem sair do console AWS.

O S3 tem um recurso muito bom para calcular o tamanho total de um caminho. Usei isso para estimar a quantidade total de dados que estava usando: modestos 9,5 TB.

Minha configuração do Luster ficou assim:

Por que fornecer Lustre? Porque reduz o tempo de download de dados em seu trabalho de treinamento do SageMaker de, nessa escala, provavelmente sólidos 90 minutos ou mais, para menos de 60 segundos.

Passei algum tempo configurando meu VPC para habilitar o treinamento no SageMaker e gravar os resultados no S3. Também estendi um AWS Deep Learning Container com todos os meus pacotes para economizar tempo durante o treinamento; em vez de pip instalar tudo para cada trabalho, faço isso uma vez ao criar a imagem e, em seguida, simplesmente faço o download dessa imagem para o meu trabalho em tempo de execução.

4/ Construa um índice de linhas json

Depois disso, gastei algumas horas tentando organizar a função pré-construída load_dataset() com meu formato webdataset. Ainda suspeito que isso seja possível, mas optei por apenas escrever meu próprio carregador de dados.

Para economizar tempo e dinheiro em meu enorme cluster de GPU, criei um índice de linhas json . Este era simplesmente um arquivo de vários GB com 50 milhões de objetos json, cada um apontando para um par válido de imagem/texto no FSx for Lustre.

Eu executei um grande trabalho do SageMaker baseado em CPU para listar todos os arquivos das minhas 18 peças no Lustre. Depois disso, carreguei, testei e criei uma função de carregador de dados muito pequena em meu notebook Studio.

5/ Executar em 192 GPUs com treinamento distribuído SageMaker

Depois disso, era hora de passar por uma única época! Aumentei meus limites para instâncias ml.p4d.24xlarge no treinamento do SageMaker em us-east-1 usando AWS Service Quota para 24. Como cada instância tem 8 GPUs, isso dá um total de 192 GPUs.

Fiz algumas execuções com apenas 1% dos meus dados gerais para garantir que o loop de treinamento funcionasse conforme o esperado; o novo recurso de piscinas quentes foi extremamente útil!

Minha configuração final de treinamento ficou assim.

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)

O que mais me surpreendeu foi que a diferença entre meus maiores e menores trabalhos, de 7 milhões para 50 milhões de pares imagem/texto, foi de apenas 20 minutos. Ambos os trabalhos usaram 24 instâncias e passaram a maior parte do tempo simplesmente inicializando as GPUs para o trabalho. Baixar a imagem, configurar o MPI, carregar os dados e preparar o NCCL levou a maior parte do tempo de execução.

Meu tempo para completar 1 época em 50 milhões de pares de imagem/texto com 200 GPUs? 15 minutos!!

Isso mesmo, 15 minutos. Das 18h30 às 18h45, a única vez que minhas GPUs foram realmente utilizadas. O trabalho geral levou cerca de uma hora para ser concluído, novamente com a maior parte do tempo gasta na inicialização, provisionamento, carregamento e upload do modelo finalizado.

Para mim, isso indica uma grande oportunidade no design e otimização de sistemas de GPU. Precisamos aceitar esse imposto de 40 minutos apenas para treinar nossos modelos? Sem chance! Vamos ser criativos sobre as soluções para corrigir isso.

Para ser claro, não treinei esse modelo para convergência. Simplesmente testei o desempenho em uma única época, que registrei em 15 minutos. Para treinar por 1000 épocas, eu orçaria 15 mil minutos, o que deve levar pouco mais de 10 dias.

7/ Implante na hospedagem do Amazon SageMaker

Depois disso, implantei o modelo básico na hospedagem do SageMaker. Como você pode esperar, um modelo treinado em uma única época não terá um desempenho muito bom, então simplesmente usei a primeira versão do Stable Diffusion em minha própria conta da AWS.

E aqui está a minha imagem resultante!

E isso é um embrulho! Nesta postagem, expliquei como trabalhei com 10 TB de dados para treinar uma única época de Stable Diffusion no treinamento distribuído do SageMaker. Meu loop de treinamento foi concluído em 15 minutos!

Espero ter começado a aguçar seu apetite por mais conteúdo sobre pré-treinamento de visão ampla e modelos de linguagem. Estou escrevendo um livro sobre esse assunto , que estará disponível em abril de 2023. Você poderá ingressar no meu Círculo de Autores em janeiro, participar da festa de lançamento, encontrar-se com outros especialistas nessa área e muito mais!