Как я тренировал 10 ТБ для Stable Diffusion на SageMaker

С подробностями о моей будущей книге по распределенному обучению с AWS!

Работать с большим количеством данных сложно, в этом нет никаких сомнений. Тем не менее, благодаря разумному выбору и простому дизайну вы также можете масштабировать свои проекты, чтобы воспроизводить их и получать самые современные результаты.

В этом посте я расскажу вам о примере, в котором я использовал AWS, в частности SageMaker, S3 и FSx для Lustre, для загрузки, обработки и обучения модели Stable Diffusion. Я использовал SageMaker Distributed Data Parallel для оптимизации градиентного спуска в масштабе. Я также использовал параллелизм заданий на модульной серверной части SageMaker для загрузки и обработки всех файлов.

Мой обучающий код основан на примере Кори Барретта из AWS MLSL, доступном здесь. Я постараюсь обновить этот репозиторий со всем исходным кодом для этого проекта к концу года.

Если вам не хватает дополнительной информации о выполнении крупномасштабных заданий в SageMaker и обо всех возможностях предварительного обучения больших графических и языковых моделей в AWS, ознакомьтесь с моей книгой, которая выйдет в апреле 2023 года!

В своей книге я посвящаю 15 главам, проводя вас через весь жизненный цикл, от поиска правильного варианта использования до подготовки вашей среды, обучения, устранения неполадок, оценки, ввода в эксплуатацию и масштабирования.

Моя распределённая учебная книга идеальна для начинающих ; это действительно для разработчиков Python от начального до среднего, которые хотят в конечном итоге получить навыки для обучения своих собственных больших концепций и языковых моделей.

Давайте погрузимся!

1/ Загрузите файлы паркета Laion-5B с заданиями SageMaker

Основным набором данных, используемым для обучения Stable Diffusion, является Laion-5B. Это набор данных с открытым исходным кодом, который предоставляет миллиарды пар изображения/текста из Интернета. Romain Beaumont, участник набора данных Laion-5B, открыл исходный код своего инструмента для загрузки и обработки этого набора данных, и я использовал его в качестве основы. Его исходный код доступен здесь. Спасибо Ромен!

Базовый набор данных состоит из двух частей: более 100 файлов паркета, которые указывают на исходные изображения и их подписи, и сами загруженные изображения.

Сначала я загрузил файлы паркета с помощью простой функции Python. К счастью, инструментарий Ромена позволяет загружать файлы непосредственно в корзину S3, что позволяет обойти любые ограничения локального хранилища. Я тестировал это несколько раз на моем экземпляре SageMaker Studio, переходя на более крупную машину, когда мне нужно было работать с большим файлом.

После этого я обернул его в один скрипт, который можно было запустить на моем задании 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/ Загрузите изображения и текстовые пары Laion-5B

После этого пришло время загрузить сами изображения. Я составил список всех файлов паркета, прошелся по нему и отправил каждый файл паркета на его собственное задание SageMaker.

Мой скрипт для загрузки пар изображения и текста выглядел так:

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

Затем я просмотрел свой список файлов паркета и отправил каждый файл паркета на свое задание . Мы называем это параллелизмом заданий . Это простой и эффективный способ запуска параллельных вычислительных заданий.

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)

Я запустил это на 18 разных заданиях, все выполнялись одновременно, в течение примерно 24 часов.

Получившиеся объекты выглядели так:

3/ Создайте том FSx для Lustre из пути S3.

Затем я потратил колоссальные двадцать минут , чтобы создать новый том FSx для Lustre из моих путей S3. Это было поразительно легко сделать; туда и обратно, включая несколько моих собственных заминок, а время создания тома составило около двадцати минут. Все это, не выходя из консоли AWS.

S3 имеет очень хорошую функцию для расчета общего размера пути. Я использовал это, чтобы оценить, сколько всего данных я использовал: скромные 9,5 ТБ.

Моя настройка Lustre выглядела так:

Зачем предоставлять Lustre? Потому что это сокращает время загрузки данных в учебном задании SageMaker с 90 или более минут в этом масштабе до менее чем 60 секунд.

Я потратил некоторое время на настройку своего VPC, чтобы включить обучение в SageMaker и запись результатов в S3. Я также расширил контейнер AWS Deep Learning Container всеми своими пакетами, чтобы сэкономить время во время обучения; вместо того, чтобы pip устанавливать все для каждого задания, я делаю это один раз при создании образа, а затем просто загружаю этот образ в свое задание во время выполнения.

4/ Создайте индекс строк json

После этого я потратил несколько часов, пытаясь сопоставить готовую функцию load_dataset() с моим форматом веб-данных. Я все еще подозреваю, что это возможно, но я решил просто написать свой собственный загрузчик данных.

Чтобы сэкономить время и деньги на моем массивном кластере графических процессоров, я создал индекс строк json . Это был просто многогигабайтный файл с 50 миллионами объектов json, каждый из которых указывал на допустимую пару изображение/текст в FSx для Lustre.

Я выполнил большое задание SageMaker на основе ЦП, чтобы перечислить все файлы из моих 18 частей в Lustre. После этого я загрузил, протестировал и создал очень маленькую функцию загрузки данных на своем ноутбуке Studio.

5/ Работа на 192 графических процессорах с распределенным обучением SageMaker.

После этого пришло время пройти через одну эпоху! Я увеличил свои лимиты для экземпляров ml.p4d.24xlarge при обучении SageMaker в регионе us-east-1 с использованием AWS Service Quota до 24. Поскольку каждый экземпляр имеет 8 графических процессоров, всего получается 192 графических процессора.

Я сделал несколько прогонов всего с 1% моих общих данных, чтобы убедиться, что тренировочный цикл работает должным образом; новая функция теплых бассейнов была очень полезной!

Мой окончательный тренировочный конфиг выглядел так.

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)

Что меня больше всего удивило, так это то, что разница между моими самыми большими и самыми маленькими заданиями, от 7 до 50 миллионов пар изображения/текста, составляла всего около 20 минут. Оба задания использовали 24 экземпляра, и оба тратили большую часть своего времени просто на инициализацию графических процессоров для задания. Загрузка образа, настройка MPI, загрузка данных и подготовка NCCL заняли большую часть времени выполнения.

Мое время завершить 1 эпоху на 50 миллионах пар изображения/текста с 200 графическими процессорами? 15 минут!!

Правильно, 15 минут. С 18:30 до 18:45 — единственный раз, когда мои графические процессоры действительно использовались. В целом работа заняла около часа, опять же, большая часть времени была потрачена на инициализацию, подготовку, загрузку и загрузку готовой модели.

Для меня это указывает на огромные возможности в разработке и оптимизации систем с графическими процессорами. Нужно ли нам принимать этот 40-минутный налог только для обучения наших моделей? Ни за что! Давайте творчески подойдем к решениям, чтобы исправить это.

Чтобы было ясно, я не обучал эту модель конвергенции. Я просто проверил производительность на одной эпохе, которую я установил на 15 минут. Чтобы обучить 1000 эпох, я бы выделил 15 тысяч минут, что должно занять чуть более 10 дней.

7/ Развертывание на хостинге Amazon SageMaker

После этого я развернул базовую модель на хостинге SageMaker. Как и следовало ожидать, модель, обученная на одной эпохе, не будет работать очень хорошо, поэтому я просто использовал первую версию Stable Diffusion в своей учетной записи AWS.

И вот мой результирующий образ!

И это обертка! В этом посте я объяснил, как я работал с 10 ТБ данных, чтобы в конечном итоге обучить одну эпоху Stable Diffusion на распределенном обучении SageMaker. Мой тренировочный цикл завершился за 15 минут!

Я надеюсь, что я начал разжигать ваш аппетит к большему количеству материалов по предварительному обучению больших моделей зрения и языка. Я пишу книгу на эту тему , которая будет доступна в апреле 2023 года. Вы сможете присоединиться к моему авторскому кругу в январе, посетить вечеринку по случаю запуска, встретиться с другими экспертами в этой области и многое другое!