Ускорение обучения PyTorch с помощью FP8

Увеличение эффективности обучения в PyTorch с использованием FP8

Как получить максимум от вашей современной видеокарты

Последние несколько лет принесли революционные достижения в области искусственного интеллекта, может быть наилучшим образом проиллюстрированные недавней популярностью и распространением приложений на основе LLM, таких как ChatGPT. Эти прорывы были порождены такими же захватывающими разработками в области оборудования, используемого для обучения моделей искусственного интеллекта. Новые и инновационные архитектуры, тензорные процессоры, специализированные аппаратные ускорители позволили достигнуть схождения искусственных интеллектуальных моделей с увеличивающимися размерами на все более быстрых скоростях. В этой статье мы сосредоточимся на одном особом преимуществе аппаратуры, специализирующейся на искусственном интеллекте – включении специализированных ядер обработки тензоров с 8-битным плавающей точкой (FP8). Появление FP8 тензорных ядер в самых современных архитектурах искусственного интеллекта (например, Nvidia Hopper, Nvidia Ada Lovelace и Habana Gaudi2) позволяет значительно увеличить количество операций с плавающей точкой в секунду (FLOPS), а также создает возможности для оптимизации памяти и энергосбережения как для обучающих, так и для выводных нагрузок искусственного интеллекта.

Для использования возможностей FP8 на уровне аппаратуры необходима соответствующая поддержка в стеке программного обеспечения и среде разработки, которую мы используем для создания обучающих и прогностических приложений искусственного интеллекта. В этой статье мы опишем, как изменить обучающий скрипт PyTorch, чтобы использовать встроенную поддержку типа данных FP8 на видеокарте Nvidia H100. Мы начнем с объяснения важности использования типа данных FP8. Затем мы рассмотрим поддерживаемые PyTorch API для работы с FP8, предоставленные библиотекой Transformer Engine и покажем, как их интегрировать в простой обучающий скрипт. Несмотря на то, что мы не будем углубляться в теорию использования FP8 для обучения искусственного интеллекта, мы отметим потенциальные проблемы, связанные с его использованием. В конце мы продемонстрируем значительные возможности оптимизации типа данных FP8.

Отказ от ответственности

Пожалуйста, не воспринимайте наше упоминание любого программного компонента, методологии или сервиса как рекомендацию к их использованию. Лучший дизайн для разработки искусственного интеллекта значительно разнится в зависимости от конкретных деталей вашей искусственной интеллектуальной задачи. Пожалуйста, также имейте в виду, что API и поведение некоторых программных пакетов и компонентов, которые мы упомянем, могут измениться к моменту вашего прочтения этой статьи. Настоятельно рекомендуется оценивать все потенциальные решения на основе самого нового аппаратного и программного обеспечения.

Мотивация

По мере развития искусственных интеллектуальных моделей требуется все более сложная аппаратура для их обучения. Видеокарта Nvidia H100, которая, как говорят, обеспечивает “беспрецедентную производительность и масштабируемость”, является (на момент написания этого) новейшим и самым мощным ускорителем искусственного интеллекта от Nvidia, специально разработанным с целью обеспечить развитие искусственного интеллекта нового поколения. С текущим ажиотажем в области искусственного интеллекта, спрос на эти видеокарты был огромным (например, здесь). Соответственно, и не удивительно, стоимость этих видеокарт была чрезвычайно высокой – возможно, даже недоступной для многих наших читателей. К счастью, провайдеры облачных услуг, такие как AWS, GCP и Microsoft Azure, предлагают доступ к машинам на основе H100 по модели “оплата за использование” (за час / за секунду), тем самым расширяя возможности использования этой технологии для гораздо большего сообщества разработчиков искусственного интеллекта.

• Использование LLM с помощью LangChain для аналитики цепи поставок – контрольная башня, работающая на основе GPT
• Визуализация данных с помощью KNIME
• Освоение обработки пакетных данных с помощью универсального инструмента для работы с данными VDK (Versatile Data Kit)

В AWS GPU H100 предлагаются в качестве компонента недавно анонсированного семейства экземпляров AWS EC2 p5. Утверждается, что эти инстансы позволят ускорить время достижения решения до 4 раз по сравнению с предыдущими генерациями GPU-инстансов EC2 и снизить стоимость обучения моделей машинного обучения на 40%.

В недавней публикации мы обсудили некоторые аспекты, которые следует учитывать при выборе инстанса для обучения моделей машинного обучения. Мы подчеркнули, что наиболее оптимальный тип инстанса будет очень сильно зависеть от конкретного проекта. В частности, когда речь идет о инстансах для обучения моделей машинного обучения, больше не всегда означает лучше. Это особенно верно для семейства p5. Да, p5, скорее всего, превзойдет другие типы инстансов — ведь H100 является безспорным монстром производительности. Но учитывая стоимость p5 ($98.32 в час за 8-графический p5.48xlarge инстанс на момент написания этого текста), вы можете найти более подходящие типы инстансов.

В следующем разделе мы обучим относительно большую модель компьютерного зрения на p5.48xlarge и сравним ее производительность с p4d.24xlarge с 8 графическими процессорами Nvidia A100.

Игрушечная модель

В приведенном ниже код-блоке мы определяем модель классификации на основе Vision Transformer (ViT) (с использованием популярной Python-библиотеки timm версии 0.9.10) вместе с случайно сгенерированным набором данных. ViT-бэкбоны имеют множество форм и размеров. Здесь мы выбрали т.н. ViT-Huge конфигурацию — с общим числом параметров 632 миллиона — чтобы лучше использовать возможности H100 для больших моделей.

import torch, timeimport torch.optimimport torch.utils.dataimport torch.distributed as distfrom torch.nn.parallel.distributed import DistributedDataParallel as DDPimport torch.multiprocessing as mp# измените размер пакета в зависимости от доступной памяти GPUbatch_size = 64from timm.models.vision_transformer import VisionTransformerfrom torch.utils.data import Dataset# используйте случайные данныеclass FakeDataset(Dataset):    def __len__(self):        return 1000000    def __getitem__(self, index):        rand_image = torch.randn([3, 224, 224], dtype=torch.float32)        label = torch.tensor(data=[index % 1000], dtype=torch.int64)        return rand_image, labeldef mp_fn(local_rank, *args):    # настройка процесса    dist.init_process_group("nccl",                            rank=local_rank,                            world_size=torch.cuda.device_count())    torch.cuda.set_device(local_rank)    device = torch.cuda.current_device()        # создание набора данных и загрузчиков данных    train_set = FakeDataset()    train_loader = torch.utils.data.DataLoader(        train_set, batch_size=batch_size,        num_workers=12, pin_memory=True)    # определение модели ViT-Huge    model = VisionTransformer(            embed_dim=1280,            depth=32,            num_heads=16,        ).cuda(device)    model = DDP(model, device_ids=[local_rank])    # определение функции потерь и оптимизатора    criterion = torch.nn.CrossEntropyLoss()    optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)    model.train()    t0 = time.perf_counter()    summ = 0    count = 0    for step, data in enumerate(train_loader):        # копирование данных на GPU        inputs = data[0].to(device=device, non_blocking=True)        label = data[1].squeeze(-1).to(device=device, non_blocking=True)          # использование смешанной точности для использования поддержки буферных значений bfloat16        with torch.autocast(device_type='cuda', dtype=torch.bfloat16):            outputs = model(inputs)            loss = criterion(outputs, label)        optimizer.zero_grad(set_to_none=True)        loss.backward()        optimizer.step()                # запоминание времени выполнения шага        batch_time = time.perf_counter() - t0        if step > 10:  # пропуск первых шагов            summ += batch_time            count += 1        t0 = time.perf_counter()        if step > 50:            break    print(f'среднее время шага: {summ/count}')if __name__ == '__main__':    mp.spawn(mp_fn,             args=(),             nprocs=torch.cuda.device_count(),             join=True)

Мы обучали эту модель на типах экземпляров p5.48xlarge и p4d.24xlarge, используя специальный контейнер AWS deep learning на основе PyTorch 2.1 AWSCdk2 CPy (763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.1.0-gpu-py310-cu121-ubuntu20.04-ec2).

Что неудивительно, производительность p5 по времени выполнения шага превосходит производительность p4d — 0,199 секунды за шаг по сравнению с 0,41 — более чем вдвое быстрее!! Это означает сокращение времени обучения больших моделей машинного обучения вдвое. Однако, если принять во внимание разницу в стоимости ($32,77 в час за p4d против $98,32 в час за p5 — на момент написания) получаем совершенно другую историю. Цена-производительность p5 примерно на 30% хуже, чем у p4d!! Это сильно расходится с основательствами 40% улучшения, объявленными в объявлении о p5.

На этом этапе у вас могут возникнуть два возможных вывода. Первая возможность заключается в том, что, несмотря на всю шумиху, p5 просто не является подходящей машиной для вас. Вторая возможность заключается в том, что p5 все же может быть пригодна, но для осуществления полного использования ее потенциала потребуются модификации вашей модели. В следующих разделах мы выберем второй подход и продемонстрируем, как использование типа экземпляра p5 уникального для данных типов экземпляров — FP8 — может полностью изменить результаты сравнительной стоимости-производительности.

Интеграция FP8 с Трасформером

Первое, на что мы должны обратить внимание, это то, что на момент написания этой статьи, PyTorch (версия 2.1) не включает нативный 8-битный тип данных с плавающей точкой. Чтобы программировать наш сценарий, используя FP8, мы будем использовать Трасформер Онжин (TE), специальную библиотеку для ускорения моделей Трасформер на графических процессорах NVIDIA. TE (версия 0.12) предустановлена в контейнере AWS PyTorch 2.1 DL.

Хотя теория использования FP8 для обучения выходит за рамки данной публикации (например, см. здесь), важно знать, что механика использования FP8 гораздо сложнее, чем у альтернатив форматов 16 бит (float16 и bfloat16). К счастью, реализация TE скрывает все вычислительные детали от пользователя. Пожалуйста, обратитесь к официальной документации, а также к этому простому примеру, чтобы узнать, как использовать API TE. Чтобы узнать больше о том, что происходит за кулисами, обязательно посмотрите следующие два видео-руководства.

Обучение FP8 с Трасформер Онжин | NVIDIA по запросу

В ходе сессии будет представлено введение в FP8 и смешанную точность, обзор особенностей Трасформер Онжин и …

www.nvidia.com

FP8 для глубокого обучения | NVIDIA по запросу

FP8 – это естественное продвижение ускорения глубокого обучения (DL) за пределы 16-битных форматов, которые обычно применяются в современных …

www.nvidia.com

Чтобы модифицировать нашу модель для использования TE, мы оборачиваем специальный трансформаторный слой TE в пользовательский класс блока трасформера, соответствующий сигнатуре блочного слоя timm.

import transformer_engine.pytorch as tefrom transformer_engine.common import recipeclass TE_Block(te.transformer.TransformerLayer):    def __init__(            self,            dim,            num_heads,            mlp_ratio=4.,            qkv_bias=False,            qk_norm=False,            proj_drop=0.,            attn_drop=0.,            init_values=None,            drop_path=0.,            act_layer=None,            norm_layer=None,            mlp_layer=None    ):        super().__init__(            hidden_size=dim,            ffn_hidden_size=int(dim * mlp_ratio),            num_attention_heads=num_heads,            hidden_dropout=proj_drop,            attention_dropout=attn_drop            )

Затем мы модифицируем инициализацию VisionTransformer для использования нашего пользовательского блочного слоя:

  model = VisionTransformer(      embed_dim=1280,      depth=32,      num_heads=16,      block_fn=TE_Block      ).cuda(device)

До сих пор мы не внесли никаких изменений, специфичных для H100 – тот же код может быть запущен на нашем экземпляре типа p4d с A100. Последнее изменение состоит в обертывании прохода прямого прохода модели с помощью контекстного менеджера te.fp8_autocast. Для этого требуется GPU, поддерживающее FP8:

with torch.autocast(device_type='cuda', dtype=torch.bfloat16):    with te.fp8_autocast(enabled=True):        outputs = model(inputs)    loss = criterion(outputs, label)

Несколько предосторожных замечаний относительно использования FP8

Использование 8-битного представления с плавающей запятой (в отличие от 16- или 32-битного представления) подразумевает меньшую точность и меньший динамический диапазон. Это может оказать существенное влияние на достижимость и/или скорость сходимости вашей модели. Несмотря на то, что имеющаяся реализация TE FP8 разработана для решения этой проблемы, нет гарантии, что она будет работать для вашей модели. Вам может потребоваться настроить механизмы FP8 (например, с использованием API TE recipe), настроить некоторые из гиперпараметров и/или ограничить применение FP8 к частям модели. Вы можете обнаружить, что несмотря на все ваши попытки, ваша модель просто несовместима с FP8.

Результаты

В таблице ниже мы подводим итоги наших экспериментов на экземплярах EC2 p4d.24xlarge и p5.48xlarge с использованием или без использования библиотеки TE. Для экспериментов на p5.48xlarge мы увеличили размер пакета в два раза, чтобы увеличить использование 80 ГБ памяти GPU. Использование FP8 позволяет уменьшить потребление памяти GPU и дополнительно увеличить размер пакета.

Мы видим, что использование блочного преобразователя TE увеличило эффективность использования ресурсов как на экземплярах p4d (~19%), так и на p5 (~32%). Использование FP8 повышает производительность на p5 примерно на 20%. Следуя оптимизациям TE и FP8, эффективность использования ресурсов экземпляра p5.48large на базе H100 превосходит эффективность использования ресурсов экземпляра p4d.24large на базе A100 – хотя и не с большим отрывом (~2%). Учитывая увеличение скорости обучения в 3 раза, мы можем с уверенностью заключить, что p5 является лучшим типом экземпляра для обучения нашей оптимизированной модели.

Обратите внимание, что относительно небольшое увеличение эффективности использования ресурсов (гораздо ниже 40%, упомянутых в объявлении p5) заставляет нас ждать дополнительных оптимизаций, специфичных для H100… но о них придется написать в следующем посте :).

Резюме

В этом посте мы продемонстрировали, как программировать сценарий обучения PyTorch, чтобы использовать 8-битные типы с плавающей точкой. Мы также показали, что использование FP8 может быть ключевым фактором для получения наилучшей производительности современных GPU, таких как Nvidia H100. Важно отметить, что возможность использования FP8, а также его влияние на производительность обучения могут сильно варьироваться в зависимости от деталей модели.

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