Comment j'ai formé 10 To pour la diffusion stable sur SageMaker

Avec des détails sur mon prochain livre sur la formation distribuée avec AWS !

Traiter beaucoup de données est difficile, cela ne fait aucun doute. Cependant, avec des choix intelligents et des conceptions simples, vous pouvez également étendre vos projets pour reproduire et produire des résultats de pointe.

Dans cet article, je vais vous présenter une étude de cas où j'ai utilisé AWS, notamment SageMaker, S3 et FSx pour Lustre, pour télécharger, traiter et former un modèle de diffusion stable. J'ai utilisé SageMaker Distributed Data Parallel pour optimiser la descente de gradient à grande échelle. J'ai également utilisé le parallélisme des tâches sur le backend modulaire de SageMaker pour télécharger et traiter tous les fichiers.

Mon code de formation est basé sur un échantillon de Corey Barrett de l'AWS MLSL, disponible ici. J'essaierai de mettre à jour ce référentiel avec tout le code source de ce projet d'ici la fin de l'année.

Si vous avez envie de plus de contenu sur l'exécution de tâches à grande échelle sur SageMaker et sur l'ensemble de la préformation de grands modèles de vision et de langage sur AWS, consultez mon livre qui sortira en avril 2023 !

Dans mon livre, je passe 15 chapitres à vous guider tout au long du cycle de vie, de la recherche du bon cas d'utilisation à la préparation de votre environnement, en passant par la formation, le dépannage, l'évaluation, l'opérationnalisation et la mise à l'échelle.

Mon livre d'entraînement distribué est idéal pour les débutants ; c'est vraiment pour les développeurs Python débutants à intermédiaires qui souhaitent éventuellement acquérir les compétences nécessaires pour former leurs propres modèles de vision et de langage étendus.

Plongeons-nous !

1/ Télécharger les fichiers parquet Laion-5B avec les jobs SageMaker

Le jeu de données de base utilisé pour former la diffusion stable est Laion-5B. Il s'agit d'un ensemble de données open source qui fournit des milliards de paires image/texte à partir d'Internet. Romain Beaumont, contributeur au jeu de données Laion-5B, a ouvert son outil pour télécharger et traiter ce jeu de données, et je l'ai utilisé comme base. Son code source est disponible ici. Merci Romain !

L'ensemble de données de base est composé de deux parties : plus de 100 fichiers de parquet qui pointent vers les images originales et leurs légendes, et les images téléchargées elles-mêmes.

Tout d'abord, j'ai téléchargé les fichiers parquet avec une simple fonction Python. Heureusement, la boîte à outils de Romain vous permet de télécharger directement dans votre compartiment S3, ce qui vous permet de contourner toutes les contraintes de stockage local. J'ai testé cela plusieurs fois sur mon instance SageMaker Studio, en passant à une machine plus grande lorsque j'avais besoin de travailler avec un fichier plus volumineux.

Après cela, je l'ai enveloppé dans un seul script que je pouvais exécuter sur mon travail 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/ Télécharger les couples image et texte Laion-5B

Après cela, il était temps de télécharger les images elles-mêmes. J'ai fait une liste de tous les fichiers de parquet, je l'ai parcourue et j'ai envoyé chaque fichier de parquet à son propre travail SageMaker.

Mon script pour télécharger les paires image et texte ressemblait à ceci :

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

Ensuite, j'ai parcouru ma liste de limes parquet et j'ai envoyé chaque lime parquet à son propre travail . Nous appelons ce travail le parallélisme , c'est un moyen simple et efficace d'exécuter des travaux de calcul parallèles.

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)

J'ai exécuté cela sur 18 tâches différentes, toutes exécutées en même temps, pendant environ 24 heures.

Les objets résultants ressemblaient à ceci :

3/ Créer un volume FSx for Lustre à partir du chemin S3

Ensuite, j'ai passé vingt minutes à créer un nouveau volume FSx for Lustre à partir de mes chemins S3. C'était incroyablement facile à faire; aller-retour incluant quelques-uns de mes propres ratés et le temps de création du volume était d'une vingtaine de minutes. Le tout sans quitter la console AWS.

S3 a une fonctionnalité très intéressante pour calculer la taille totale d'un chemin. J'ai utilisé cela pour estimer la quantité totale de données que j'utilisais : un modeste 9,5 To.

Ma configuration Lustre ressemblait à ceci:

Pourquoi fournir Lustre ? Parce qu'il réduit le temps de téléchargement de vos données sur votre travail de formation SageMaker de, à cette échelle, probablement 90 minutes ou plus, à moins de 60 secondes.

J'ai passé du temps à configurer mon VPC pour permettre la formation sur SageMaker et écrire les résultats sur S3. J'ai également étendu un AWS Deep Learning Container avec tous mes packages pour gagner du temps pendant la formation ; au lieu de tout installer par pip pour chaque travail, je le fais une fois lors de la création de l'image, puis je télécharge simplement cette image sur mon travail au moment de l'exécution.

4/ Construire un index de lignes json

Après cela, j'ai passé quelques heures à essayer de démêler la fonction prédéfinie load_dataset() avec mon format webdataset. Je soupçonne toujours que cela est possible, mais j'ai choisi d'écrire simplement mon propre chargeur de données.

Pour gagner du temps et de l'argent sur mon énorme cluster GPU, j'ai construit un index de lignes json . Il s'agissait simplement d'un fichier de plusieurs Go avec 50 millions d'objets json, chacun pointant vers une paire image/texte valide sur FSx pour Lustre.

J'ai exécuté un gros travail SageMaker basé sur le processeur pour répertorier tous les fichiers de mes 18 parties sur Lustre. Après cela, j'ai chargé, testé et construit une très petite fonction de chargeur de données sur mon ordinateur portable Studio.

5/ Exécuter sur 192 GPU avec la formation distribuée SageMaker

Après cela, il était temps de traverser une seule époque ! J'ai augmenté mes limites pour les instances ml.p4d.24xlarge sur la formation SageMaker dans us-east-1 en utilisant AWS Service Quota à 24. Étant donné que chaque instance a 8 GPU, cela fait un total de 192 GPU.

J'ai effectué quelques courses avec seulement 1 % de mes données globales pour m'assurer que la boucle d'entraînement fonctionnait comme prévu ; la nouvelle fonctionnalité de piscines chaudes a été extrêmement utile !

Ma configuration d'entraînement finale ressemblait à ceci.

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)

Ce qui m'a le plus surpris, c'est que la différence entre mes travaux les plus gros et les plus petits, de 7 millions à 50 millions de paires image/texte, n'était que d'environ 20 minutes. Les deux tâches utilisaient 24 instances et passaient la majorité de leur temps à simplement initialiser les GPU pour la tâche. Le téléchargement de l'image, la configuration de MPI, le chargement des données et la préparation de NCCL ont pris la grande majorité du temps d'exécution.

Mon temps pour terminer 1 époque sur 50 millions de paires image/texte avec 200 GPU ? 15 minutes!!

C'est vrai, 15 minutes. De 18h30 à 18h45, la seule fois où mes GPU ont été réellement utilisés. Le travail global a duré environ une heure, encore une fois, la majorité du temps étant consacrée à l'initialisation, à l'approvisionnement, au chargement et au téléchargement du modèle fini.

Pour moi, cela indique une énorme opportunité dans la conception et l'optimisation des systèmes GPU. Avons-nous besoin d'accepter cette taxe de 40 minutes juste pour former nos modèles ? Certainement pas! Soyons créatifs sur les solutions pour résoudre ce problème.

Pour être clair, je n'ai pas formé ce modèle à la convergence. J'ai simplement testé les performances sur une seule époque, que j'ai chronométrée à 15 minutes. Pour m'entraîner pendant 1000 époques, je prévoirais 15 000 minutes, ce qui devrait prendre un peu plus de 10 jours.

7/ Déployer sur l'hébergement Amazon SageMaker

Après cela, j'ai déployé le modèle de base sur l'hébergement SageMaker. Comme vous vous en doutez, un modèle formé sur une seule époque ne fonctionnera pas très bien, j'ai donc simplement utilisé la première version de Stable Diffusion dans mon propre compte AWS.

Et voici mon image résultante!

Et c'est un enveloppement! Dans cet article, j'ai expliqué comment j'ai travaillé avec 10 To de données pour finalement former une seule époque de diffusion stable sur la formation distribuée SageMaker. Ma boucle d'entraînement bouclée en 15 minutes !

J'espère que j'ai commencé à aiguiser votre appétit pour plus de contenu sur la pré-formation des modèles de grande vision et de langage. J'écris un livre sur ce sujet , qui sera disponible en avril 2023. Vous pourrez rejoindre mon cercle d'auteurs en janvier, assister à la soirée de lancement, rencontrer d'autres experts du domaine dans ce domaine, et plus encore !