Высокопроизводительная передача памяти GPU на AWS Sagemaker Hyperpod

Современная инфраструктура глубокого обучения зачастую требует передачи больших объемов данных между графическими процессорами (GPU) разных машин. В Perplexity мы столкнулись с уникальной технической задачей: эффективная передача несмежных областей памяти GPU между машинами с максимально возможной скоростью. Наша целевая платформа, экземпляры AWS p5, обеспечивает впечатляющую пропускную способность сети в 3200 Гбит/с через 32 сетевые карты. В этой статье мы расскажем о нашем опыте создания специализированного высокопроизводительного сетевого решения, достигающего 97,1% от этой теоретической пропускной способности.

Техническая задача

Наш сценарий использования предполагал несколько ключевых технических требований:

• Высокоскоростная передача несмежных блоков памяти между удаленными GPU
• Возможность динамического добавления или удаления узлов из развертываний Kubernetes без нарушения текущих операций
• Поддержка шаблонов взаимодействия между равноправными узлами (peer-to-peer)

Хотя библиотека NVIDIA NCCL является де-факто стандартом для распределенного глубокого обучения, она не идеально подходила для нашего случая:

• NCCL отлично справляется с коллективным взаимодействием, но требует создания статического «мира», что влечет за собой перезапуск всего кластера при изменении количества участвующих узлов.
• Модель синхронного взаимодействия NCCL усложняет нашу асинхронную рабочую нагрузку
• Мы хотели получить прямой контроль над нашими шаблонами передачи памяти для оптимизации
• Создание собственного решения предоставило ценные возможности для обучения

Современные высокопроизводительные сети

Чтобы понять наше решение, давайте сначала разберемся, чем современные высокопроизводительные сети отличаются от традиционных.

Большинство сетей, которые мы используем ежедневно, опираются на протоколы TCP/IP, где приложения взаимодействуют с сетевой картой через ядро операционной системы с использованием сокетов. Однако высокопроизводительные сети используют RDMA (Remote Direct Memory Access) — принципиально иной стек аппаратного и программного обеспечения, обеспечивающий прямой доступ к памяти между машинами без участия центрального процессора.

AWS предоставляет Elastic Fabric Adapter (EFA) — специализированный сетевой интерфейс, реализующий проприетарный протокол Amazon под названием Scalable Reliable Datagram (SRD). В отличие от традиционных сетей TCP/IP, где данные должны многократно копироваться между пользовательским пространством, пространством ядра и сетевыми буферами, EFA с поддержкой RDMA позволяет осуществлять прямую передачу данных между памятью GPU и сетевой картой, полностью минуя центральный процессор.

Философия проектирования высокопроизводительных сетей

Построение высокопроизводительных сетевых систем требует пересмотра ряда фундаментальных предположений:

1. Владение буфером: в отличие от традиционных сокетов, где ядро управляет сетевыми буферами и требует копирования данных между пользовательским пространством и пространством ядра, RDMA требует, чтобы приложения сами управляли своими буферами. Когда приложение инициирует сетевую операцию, оно передает право владения буфером сетевой карте до завершения операции, что исключает необходимость копирования данных.
2. Регистрация памяти: приложения должны регистрировать области памяти в ядре операционной системы. Ядро настраивает отображения виртуальных адресов, позволяющие центральному процессору, GPU и сетевым картам «понимать» одни и те же виртуальные адреса. Эта регистрация является разовой операцией, обеспечивающей последующую передачу данных без копирования.
3. Плоскость управления против плоскости данных: высокопроизводительные сети разделяют операции на две категории:
   1. Операции плоскости управления (такие как настройка соединения и регистрация памяти) проходят через ядро для обеспечения безопасности
   2. Операции плоскости данных (непосредственная передача данных) обходят ядро для обеспечения максимальной производительности
4. Прием перед передачей: без буферов, управляемых ядром, приложения должны заранее опубликовать операции приема, указывая место, куда должны быть помещены входящие данные. Это фундаментальное отличие от модели сокетов, где приложения могут принимать данные в любое время.
5. Опрос завершения операций: вместо ожидания уведомлений от ядра через такие механизмы, как epoll, приложения напрямую опрашивают аппаратные очереди завершения. Это исключает накладные расходы на системные вызовы и позволяет мгновенно реагировать на выполненные операции.
6. Учет аппаратной топологии: понимание и оптимизация под аппаратную топологию имеют решающее значение для достижения максимальной производительности.

Понимание аппаратной топологии

Экземпляры AWS p5 обладают сложной внутренней архитектурой. Как показано ниже, каждый экземпляр содержит два процессорных разъема, формирующих два узла NUMA, при этом каждый узел NUMA подключен к четырем коммутаторам PCIe:

Развернуть

Под каждым коммутатором PCIe находятся:

• Четыре сетевые карты EFA со скоростью 100 Гбит/с
• Один GPU NVIDIA H100
• Один NVMe SSD

Пути передачи данных для TCP/IP и RDMA демонстрируют принципиальные различия в подходах:

Развернуть

При использовании TCP/IP (слева) данные должны копироваться несколько раз:

1. Из памяти GPU в буфер приложения в пользовательском пространстве
2. Из буфера в пользовательском пространстве в буфер ядра
3. Из буфера ядра на сетевую карту
4. На принимающей стороне происходит обратный процесс

Каждая операция копирования потребляет циклы процессора и пропускную способность памяти. Приложение также должно выполнять переключение контекста между пользовательским пространством и пространством ядра для каждой сетевой операции.

В отличие от этого, RDMA (справа) обеспечивает передачу данных без копирования:

1. Сетевая карта считывает данные непосредственно из памяти GPU
2. Данные передаются напрямую на удаленную сетевую карту
3. Удаленная сетевая карта записывает данные непосредственно в память целевого GPU
4. Приложению нужно лишь проверить очередь завершения (CQ) в пользовательском пространстве, чтобы узнать о завершении передачи

При использовании RDMA и надлежащем сопряжении оборудования для передачи данных между двумя GPU требуется только прохождение через локальный коммутатор PCIe и сеть:

Развернуть

Напротив, при передаче по протоколу TCP/IP данные должны многократно копироваться через основную память, что вызывает значительную перегрузку шины PCIe:

Развернуть

Создание с использованием libfabric

Мы использовали libfabric — платформу, предоставляющую универсальный интерфейс для сетевых служб. Наша реализация использует два типа операций RDMA:

• Двусторонний RDMA (SEND/RECV) для управляющих сообщений, содержащих метаданные об областях памяти (например, смещения и размеры).
• Односторонний RDMA WRITE для фактической передачи данных, где каждая операция WRITE обрабатывает один непрерывный блок памяти.

Путь к пиковой производительности

Наша разработка прошла несколько этапов:

1. Реализована базовая однонаправленная передача сообщений с использованием SEND/RECV
2. Расширено до двунаправленного взаимодействия
3. Добавлена поддержка GPUDirect RDMA WRITE для прямой передачи GPU-GPU
4. Расширено для обработки нескольких параллельных передач
5. Внедрена очередь операций для обеспечения отказоустойчивости
6. Достигнуто использование пропускной способности 97,4% на одной сетевой карте

При масштабировании до 32 сетевых карт мы реализовали несколько важнейших оптимизаций:

1. Очередь операций: вместо прямой отправки операций на сетевые карты мы используем очередь на уровне приложения. Это обеспечивает устойчивость к перегрузкам сети и упрощает модель программирования.
2. Предварительный прогрев сети: заблаговременное установление соединений улучшило производительность при запуске.
3. Многопоточность: выделенные треды для сетевых операций каждого GPU.
4. Привязка тредов к ядрам CPU: связывание тредов с конкретными ядрами центрального процессора для предотвращения эффектов NUMA и промахов кэша.
5. Шардирование состояния: снижение конкуренции между тредами.
6. Пакетная обработка операций: одновременная отправка нескольких операций.
7. Отложенная публикация операций: операции сначала ставятся в очередь в приложении. После опроса очереди завершения мы пытаемся отправить ожидающие операции на сетевую карту, обеспечивая эффективное использование сетевых ресурсов.
8. Распределение ресурсов с учетом NUMA: выделение ресурсов libfabric, таких как очереди завершения, на правильном узле NUMA для минимизации задержек доступа к памяти.

Благодаря этим оптимизациям мы достигли итоговой производительности 3108 Гбит/с — 97,1% от теоретической максимальной пропускной способности.

На видео ниже показана наша программа командной строки в действии. Она передает несмежные блоки памяти GPU на удаленный узел, достигая скорости передачи 3108,283 Гбит/с, что демонстрирует практически теоретическое использование пропускной способности сетевой инфраструктуры:

Заключение

Создание высокопроизводительной сетевой системы требует понимания как аппаратной архитектуры, так и принципов системного проектирования. Хотя такие библиотеки, как NCCL, предоставляют отличные решения для распространенных шаблонов, иногда для конкретных требований необходимы специализированные решения.

Наш опыт показывает, что достижение теоретически возможной производительности сети достижимо при тщательном подходе к системной архитектуре, аппаратной топологии и применении различных методов оптимизации. Ключ заключается не только в понимании отдельных компонентов, но и в том, как они взаимодействуют, формируя целостную систему.

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

Цитата

@online{perplexity2025HighPerformanceGpuMemoryTransferOnAwsSagemakerHyperpod,
  author = {{Perplexity} and Lequn Chen and Phi Hoang},
  title = {{High-performance GPU memory transfer on AWS Sagemaker Hyperpod}},
  date = {2025-02-10},
  url = {https://www.perplexity.ai/hub/blog/high-performance-gpu-memory-transfer-on-aws-sagemaker-hyperpod}
}