Подробный анализ узких мест в передаче данных, их выявление и устранение с помощью систем NVIDIA Nsight™ – часть 3
Делиться
Это третья часть серии статей об оптимизации передачи данных с помощью профилировщика NVIDIA Nsight™ Systems (nsys). Первая часть была посвящена копированию данных между ЦП и ГП, а вторая — между ГП и ЦП. В этой статье мы перейдем к рассмотрению передачи данных между ГП.
В настоящее время довольно часто обучение моделей ИИ/машинного обучения — особенно больших моделей — распределяется между несколькими графическими процессорами (GPU). Хотя существует множество различных схем для такого распределения, все они имеют общую черту — зависимость от постоянной передачи данных — таких как градиенты, веса, статистика и/или метрики — между GPU на протяжении всего обучения. Как и в случае с другими типами передачи данных, которые мы анализировали в наших предыдущих статьях, здесь также некачественная реализация может легко привести к неэффективному использованию вычислительных ресурсов и неоправданному завышению стоимости обучения. Оптимизация связи между GPU — это активная область исследований и инноваций, включающая как аппаратную, так и программную разработку.
В этой статье мы сосредоточимся на наиболее распространенной форме распределенного обучения — обучении с распределением данных. При обучении с распределением данных идентичные копии модели машинного обучения хранятся на каждом графическом процессоре. Каждый пакет входных данных равномерно распределяется между графическими процессорами, каждый из которых выполняет шаг обучения для вычисления локальных градиентов. Затем локальные градиенты передаются и усредняются по всем графическим процессорам, в результате чего для каждой из копий модели выполняется идентичное обновление градиента. Используя профилировщик NVIDIA Nsight™ Systems (nsys), мы проанализируем влияние передачи градиентов между графическими процессорами на производительность обучения тестовой модели и оценим несколько методов снижения накладных расходов.
Отказ от ответственности
Приведенный в этом посте код предназначен для демонстрационных целей; пожалуйста, не полагайтесь на его точность или оптимальность. Пожалуйста, не воспринимайте упоминание нами какого-либо инструмента, фреймворка, библиотеки, сервиса или платформы как одобрение их использования.
Благодарим Ицхака Леви за его вклад в эту публикацию.
Выбор экземпляра для распределенного обучения
В предыдущей статье «Выбор экземпляра для глубокого обучения» мы обсуждали важность выбора типа экземпляра, наиболее подходящего для вашей рабочей нагрузки в области ИИ/машинного обучения, и потенциальное влияние этого выбора на успех вашего проекта. При выборе типа экземпляра для рабочей нагрузки, включающей большой объем трафика между графическими процессорами, следует обратить внимание на то, как осуществляется такая связь, включая: топологию экземпляра, межпроцессорные соединения графических процессоров, максимальную пропускную способность и задержку.
В этом посте мы ограничимся обсуждением распределения нагрузки между графическими процессорами на одном экземпляре. Мы поэкспериментируем с двумя типами экземпляров: Amazon EC2 g6e.48xlarge с 8 графическими процессорами NVIDIA L40S и Amazon EC2 p4d.24xlarge с 8 графическими процессорами NVIDIA A100. На каждом из них будет запущен образ AMI AWS Deep Learning (Ubuntu 24.04) с PyTorch (2.8), профилировщиком nsys-cli (версия 2025.6.1) и библиотекой NVIDIA Tools Extension (NVTX).
Одно из главных различий между двумя типами экземпляров заключается в способе подключения графических процессоров: в g6e.48xlarge связь между графическими процессорами осуществляется по PCI Express (PCIe), тогда как p4d.24xlarge включает NVIDIA NVLink™ — выделенное оборудование для обеспечения высокоскоростной связи между графическими процессорами. Скорость передачи данных по шине PCIe значительно ниже, чем по NVLink. Хотя этого может быть достаточно для рабочих нагрузок с низким соотношением скорости передачи данных к вычислительным ресурсам, это может существенно снизить производительность для рабочих нагрузок с высокой скоростью передачи данных.
Для определения топологии типов экземпляров выполните следующую команду:
nvidia-smi topo -m
На устройстве g6e.48xlarge мы получаем следующие результаты:
Графические процессоры на одном узле NUMA соединены каналом «NODE», а графические процессоры на разных узлах NUMA — каналом «SYS». Оба канала проходят через интерфейс PCIe, а также через другой интерфейс; ни один из них не является прямым соединением (так называемый канал «PIX»). Позже мы увидим, как это может повлиять на производительность.
На плате p4d.24xlarge каждая пара графических процессоров соединена выделенным каналом NVLink:
Игрушечная модель
Для облегчения обсуждения мы создадим простой эксперимент по распределенному обучению на основе данных.
Мы выбираем модель с относительно высоким соотношением затрат на передачу данных к вычислительным ресурсам — модель классификации изображений Vision Transformer (ViT)-L/32 с примерно 306 миллионами параметров — и синтетический набор данных, который мы будем использовать для ее обучения:
import os, time, torch, nvtx import torch.nn as nn import torch.optim as optim import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data.distributed import DistributedSampler from torch.utils.data import Dataset, DataLoader from torchvision.models import vit_l_32 WORLD_SIZE = int(os.environ.get(«WORLD_SIZE», 1)) BATCH_SIZE = 32 IMG_SIZE = 224 WARMUP_STEPS = 10 PROFILE_STEPS = 3 COOLDOWN_STEPS = 1 TOTAL_STEPS = WARMUP_STEPS + PROFILE_STEPS + COOLDOWN_STEPS N_WORKERS = 8 def get_model(): return vit_l_32(weights=None) # Синтетическая модель Класс FakeDataset(Dataset): def __len__(self): return TOTAL_STEPS * BATCH_SIZE * WORLD_SIZE def __getitem__(self, index): img = torch.randn((3, IMG_SIZE, IMG_SIZE)) label = torch.randint(0, 1000, (1,)).item() return img, label def get_data_iter(rank, world_size): dataset = FakeDataset() sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True) train_loader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=sampler, num_workers=N_WORKERS, pin_memory=True) return iter(train_loader)
Мы определяем вспомогательную функцию, которую будем использовать для настройки обучения с помощью PyTorch DistributedDataParallel (DDP). Мы настраиваем пакет распределенной связи PyTorch для использования библиотеки NVIDIA Collective Communications Library (NCCL) в качестве бэкэнда связи и оборачиваем модель в контейнер DistributedDataParallel.
def configure_ddp(model, rank): dist.init_process_group(«nccl») model = DDP(model, device_ids=[rank], bucket_cap_mb=2000) return model
Обратите внимание, что мы наивно устанавливаем емкость буфера DDP на 2 ГБ. DDP группирует градиенты модели в буферы и выполняет уменьшение градиента для каждого буфера в отдельной команде. Установка емкости буфера на 2 ГБ означает, что все ~306 миллионов (FP32) градиентов поместятся в один буфер (306 миллионов x 4 байта на градиент = ~1,22 ГБ), и уменьшение градиента произойдет только после того, как все градиенты будут вычислены.
Как и в наших предыдущих постах, мы программно планируем запуск профилировщика nsys и обозначим различные этапы этапа обучения цветовыми аннотациями NVTX:
def train(use_ddp=False): # Определяем переменные окружения, установленные torchrun rank = int(os.environ.get(«RANK», 0)) local_rank = int(os.environ.get(«LOCAL_RANK», 0)) torch.cuda.set_device(local_rank) model = get_model().to(local_rank) criterion = nn.CrossEntropyLoss().to(local_rank) if use_ddp: model = configure_ddp(model, rank) optimizer = optim.SGD(model.parameters()) data_iter = get_data_iter(rank, WORLD_SIZE) model.train() for i in range(TOTAL_STEPS): # Планируем профилирование if i == WARMUP_STEPS: torch.cuda.synchronize() start_time = time.perf_counter() torch.cuda.profiler.start() elif i == WARMUP_STEPS + PROFILE_STEPS: torch.cuda.synchronize() torch.cuda.profiler.stop() end_time = time.perf_counter() with nvtx.annotate(f»Batch {i}», color=»blue»): with nvtx.annotate(«get batch», color=»red»): data, target = next(data_iter) data = data.to(local_rank, non_blocking=True) target = target.to(local_rank, non_blocking=True) with nvtx.annotate(«forward», color=»green»): output = model(data) loss = criterion(output, target) with nvtx.annotate(«backward», color=»purple»): optimizer.zero_grad() loss.backward() with nvtx.annotate(«optimizer step», color=»yellow»): optimizer.step() if use_ddp: dist.destroy_process_group() if rank == 0: total_time = end_time — start_time print(f»Throughput: {PROFILE_STEPS/total_time:.2f} steps/sec») if __name__ == «__main__»: # включить ddp при запуске с torchrun train(use_ddp=»RANK» in os.environ)
Мы установили переменную среды NCCL_DEBUG, чтобы иметь возможность отслеживать, как NCCL настраивает каналы связи между графическими процессорами.
export NCCL_DEBUG=INFO
Производительность одной видеокарты
Мы начинаем наши эксперименты с запуска скрипта на одном графическом процессоре без поддержки DDP:
nsys profile —capture-range=cudaProfilerApi —capture-range-end=stop —trace=cuda,nvtx,osrt,nccl —output=baseline python train.py
Обратите внимание на включение nccl в раздел трассировки; это будет иметь решающее значение для анализа обмена данными между графическими процессорами в экспериментах с использованием нескольких графических процессоров.
Пропускная способность базового эксперимента составляет 8,91 шагов в секунду на графическом процессоре L40S и 5,17 шагов в секунду на A100: более новая архитектура NVIDIA Ada Lovelace показывает лучшие результаты, чем архитектура NVIDIA Ampere, на нашей тестовой модели. На изображении ниже показана временная шкала для эксперимента с L40S.
В этом посте мы сосредоточимся на части временной шкалы, относящейся к CUDA. Рассматривая строку NVTX, мы видим повторяющийся красный (копирование с ЦП на ГП), зеленый (прямой проход), фиолетовый (обратный проход) и желтый (обновление оптимизатора) паттерн, составляющий наш шаг обучения. Однако следует отметить, что фиолетовый цвет выглядит как крошечный пик, в то время как на практике обратный проход заполняет весь промежуток между прямым проходом и блоком оптимизатора: похоже, библиотека NVTX зафиксировала только первоначальный запуск графа автоградуировки.
Мы будем использовать эту кривую в качестве сравнительного базового уровня для наших следующих экспериментов.
DDP с 1 графическим процессором
Мы оцениваем влияние оболочки DDP, запустив torchrun на одном графическом процессоре:
torchrun —nproc_per_node=1 —no-python nsys profile —capture-range=cudaProfilerApi —capture-range-end=stop —trace=cuda,nvtx,nccl,osrt —output=ddp-1gpu python train.py
В результате пропускная способность падает до 8,40 шагов в секунду на графическом процессоре L40S и до 5,04 шагов в секунду на A100. Даже при отсутствии межпроцессорной связи DDP вносит накладные расходы, которые могут снизить пропускную способность примерно на 3–7%. Важный вывод из этого заключается в том, что всегда следует убедиться в том, что обучение на одном графическом процессоре выполняется без DDP.
Трассировка nsys эксперимента L40S подтверждает наличие избыточных данных DDP:
Главное отличие от этапа обучения в предыдущем примере заключается в большом объеме копирования данных между устройствами (DtoD) в конце обратного блока и непосредственно перед блоком оптимизатора (выделено в приведенном выше примере). Это и есть работа DDP: даже при отсутствии межпроцессорного сокращения градиентов, DDP подготавливает локальные градиенты в выделенном блоке памяти для сокращения, а затем копирует результаты обратно в свойство grad каждого параметра в рамках подготовки к обновлению градиента. (Обратите внимание, что копирование DtoD происходит между двумя ячейками памяти на одном и том же графическом процессоре, а не между двумя разными графическими процессорами).
DDP с несколькими графическими процессорами
Далее мы оценим влияние совместного использования градиентов между 2, 4 и 8 графическими процессорами:
torchrun —nproc_per_node=8 —no-python nsys profile —capture-range=cudaProfilerApi —capture-range-end=stop —trace=cuda,nvtx,nccl,osrt —output=ddp-8gpu_%q{RANK} python train.py
В таблице ниже показано влияние на пропускную способность обучения на устройствах g6e.48xlarge и p4d.24xlarge:
Производительность обучения DDP резко падает на экземпляре g6e.48xlarge: пропускная способность 8-GPU более чем в 6 раз ниже, чем у одного GPU. В логах NCCL содержится несколько строк, описывающих пути связи между GPU, например:
NCCL INFO Канал 00 : 0[0] -> 1[1] через SHM/direct/direct
Это означает, что передача данных между каждой парой графических процессоров осуществляется через общую память центрального процессора, что значительно ограничивает пропускную способность канала связи.
На платформе p4d.24xlarge, где соединения NVLink обеспечивают прямую связь между двумя точками (P2P), пропускная способность 8-GPU всего на 8% ниже базового результата:
NCCL INFO Канал 00/0 : 0[0] -> 1[1] через P2P/CUMEM/read
Несмотря на то, что каждый отдельный L40S превосходит A100 на нашей тестовой модели на 72%, наличие NVLink делает p4d.24xlarge более оптимальным, чем g6e.48xlarge, для работы DDP на 8 графических процессорах.
Проблема с пропускной способностью передачи данных на g6e.48xlarge легко видна в трассировке nsys. Здесь мы используем опцию «множественный просмотр», чтобы отобразить активность всех процессов DDP на одной временной шкале:
Уменьшение градиента происходит в вызове ncclAllReduce для строки NCCL. Это происходит сразу после завершения локального вычисления градиента и непосредственно перед операцией DtoD с памятью, описанной выше, которая копирует уменьшенные градиенты обратно в поле grad каждого параметра. На g6e.48xlarge операция NCCL доминирует в обратном проходе, составляя приблизительно 84% от общего времени шага (587 из 701 миллисекунды).
На модели p4d.24xlarge вызов NCCL занимает гораздо меньшую часть этапа обучения:
В следующих разделах мы обсудим несколько методов оптимизации DDP и оценим их влияние на производительность во время выполнения. Мы ограничимся оптимизациями на уровне PyTorch, оставив другие методы (например, настройку NCCL/ОС/сети) для другой публикации.
Общий подход к решению проблемы узких мест в обмене данными заключается в уменьшении частоты обмена данными. В рабочих нагрузках DDP этого можно достичь путем увеличения размера пакета данных (т.е. увеличения вычислительной нагрузки на шаг) или, если объем памяти этого не позволяет, путем применения накопления градиентов — вместо применения редукции градиентов и обновления на каждом шаге, накапливать локальные градиенты в течение N шагов и применять редукцию на каждом N-м шаге. Оба метода увеличивают эффективный размер пакета — общее количество выборок между обновлениями градиентов, что может повлиять на сходимость модели и потребовать настройки параметров оптимизации. В нашем обсуждении мы предполагаем, что эффективный размер пакета фиксирован.
Мы рассмотрим четыре метода. Первые два направлены на снижение накладных расходов DDP. Последние два напрямую решают проблему узкого места в передаче данных.
Оптимизация 1: Объявление статического графа
Первое изменение — это передача параметра static_graph=True контейнеру DDP. В общем, модели могут содержать параметры, для которых градиенты не вычисляются на каждом шаге — в документации DDP это называется «неиспользуемыми параметрами». Это часто встречается в моделях с условной логикой. DDP включает специальные механизмы для идентификации и обработки неиспользуемых параметров. В случае нашей тестовой модели все градиенты вычисляются на каждом шаге — наш граф является «статическим». Явное объявление графа статическим снижает накладные расходы, связанные с обработкой динамического набора градиентов.
def configure_ddp(model, rank): dist.init_process_group(«nccl») model = DDP(model, device_ids=[rank], static_graph=True, bucket_cap_mb=2000) return model
В случае нашей модельной модели влияние этого изменения незначительно как для g6e.48xlarge, так и для p4d.24xlarge. Без дальнейших промедлений переходим к следующей оптимизации.
Оптимизация 2: Повышение эффективности использования памяти
Второй метод решает проблему значительного копирования памяти DtoD, выявленную выше. Вместо копирования градиентов в блоки памяти связи NCCL и обратно, мы можем явно установить параметр gradients таким образом, чтобы он указывал непосредственно на буферы связи NCCL. Следовательно, одна и та же память используется для хранения локальных градиентов, выполнения редукции градиентов и применения обновлений градиентов. Это настраивается путем передачи параметра gradient_as_bucket_view=True контейнеру DDP:
def configure_ddp(model, rank): dist.init_process_group(«nccl») model = DDP(model, device_ids=[rank], static_graph=True, gradient_as_bucket_view=True, bucket_cap_mb=2000) return model
На приведенном ниже графике, полученном с помощью p4d.24xlarge, мы больше не видим блок копирования памяти DtoD между этапами all-reduce и (желтым) оптимизатором:
В нашем примере с игрушечным программным обеспечением эта оптимизация повышает производительность всего на 1%.
Оптимизация 3: Градиентное сжатие
Распространенный подход к решению проблем с пропускной способностью связи заключается в применении алгоритмов сжатия для уменьшения размера полезной нагрузки. PyTorch DDP предоставляет ряд специализированных коммуникационных механизмов, которые автоматизируют сжатие градиентов перед редукцией NCCL и последующую декомпрессию. В приведенном ниже блоке кода мы применяем сжатие градиентов типа bfloat16:
import torch.distributed.algorithms.ddp_comm_hooks.default_hooks as ddp_hks def configure_ddp(model, rank): dist.init_process_group(«nccl») model = DDP(model, device_ids=[rank], static_graph=True, gradient_as_bucket_view=True, bucket_cap_mb=2000) model.register_comm_hook(state=None, hook=ddp_hks.bf16_compress_hook) return model
На сильно перегруженном экземпляре g6e.48xlarge сжатие bfloat16 приводит к значительному ускорению на 65%! Трассировка nsys показывает как вызов NCCL уменьшенного размера, так и новые операции сжатия:
На платформе p4d.24xlarge накладные расходы на операции сжатия перевешивают выигрыш в скорости передачи данных, что приводит к общему снижению пропускной способности:
PyTorch предлагает более агрессивный алгоритм сжатия, чем bfloat16, называемый PowerSGD. Ниже мы приводим простой пример использования — на практике это может потребовать значительной настройки. Подробности см. в документации:
from torch.distributed.algorithms.ddp_comm_hooks import powerSGD_hook def configure_ddp(model, rank): dist.init_process_group(«nccl») model = DDP(model, device_ids=[rank], static_graph=True, gradient_as_bucket_view=True, bucket_cap_mb=2000) state = powerSGD_hook.PowerSGDState( process_group=None ) model.register_comm_hook( state, hook=powerSGD_hook.powerSGD_hook ) return model
PowerSGD оказывает существенное влияние на экземпляр g6e.48xlarge, увеличивая пропускную способность до 7,5 шагов в секунду — более чем в пять раз быстрее, чем базовый результат! Обратите внимание на уменьшение размера ядра NCCL в полученном трассировочном файле:
Важно отметить, что эти алгоритмы сжатия являются алгоритмами с потерями точности и должны использоваться с осторожностью, поскольку они могут повлиять на сходимость вашей модели.
Оптимизация 4: Параллельная обработка градиентного сокращения
DDP группирует параметры в несколько сегментов и запускает градиентное сокращение для каждого сегмента независимо — как только сегмент заполняется. Это позволяет выполнять градиентное сокращение заполненных сегментов, пока продолжается вычисление градиента (других сегментов). Степень перекрытия зависит от количества сегментов DDP, которое мы контролируем с помощью параметра bucket_cap_mb контейнера DDP. Напомним, что в нашей первоначальной реализации мы явно установили это значение равным 2000 (2 ГБ), что (учитывая размер нашей модели) соответствовало одному сегменту DDP. В таблице ниже показана пропускная способность для различных значений bucket_cap_mb. Оптимальная настройка будет варьироваться в зависимости от особенностей модели и среды выполнения.
Обратите внимание на существенное улучшение на 12% при использовании нескольких сегментов на g6e.48xlarge с сжатием BF16. PowerSGD, с другой стороны, работает лучше всего при применении к одному сегменту.
На изображении ниже, полученном на устройстве p4d.24xlarge с параметром bucket_cap_mb, установленным на 100, мы можем увидеть влияние этой оптимизации на трассировку профилировщика:
Вместо единой операции NCCL all-reduce теперь у нас есть 11 меньших блоков, работающих параллельно с локальным вычислением градиента.
Результаты
Результаты наших экспериментов суммированы в следующей таблице:
На p4d.24xlarge, где данные передаются по NVLink, общий эффект составил лишь скромные 4% ускорения. Но на g6e.48xlarge прирост был значительным и в основном обусловлен градиентным сжатием — 86% прироста для схемы BF16 и более чем 5-кратного улучшения для (простой реализации) PowerSGD.
Важно отметить, что эти результаты могут значительно различаться в зависимости от архитектуры модели и среды выполнения.
Краткое содержание
На этом завершается наша серия из трех статей, посвященных выявлению и устранению узких мест в передаче данных с помощью профилировщика NVIDIA Nsight™ Systems (nsys). В каждой из статей мы демонстрировали тенденцию передачи данных к возникновению узких мест в производительности и неэффективному использованию ресурсов. В каждом случае мы использовали профилировщик nsys для выявления узких мест и оценки влияния различных методов оптимизации.
Достигнутые нами в каждом из изученных нами сценариев ускорения подтверждают важность интеграции регулярного использования таких инструментов, как nsys profiler, в рабочий процесс разработки ИИ/машинного обучения и подчеркивают возможность — даже для специалистов, не работающих с CUDA, — добиться значительного повышения производительности и снижения затрат на ИИ/машинное обучение.
Источник: towardsdatascience.com
