Високопродуктивна передача пам'яті GPU на AWS Sagemaker Hyperpod

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

Технічний виклик

Наш сценарій використання висунув кілька ключових технічних вимог:

• Високошвидкісна передача між віддаленими GPU фрагментів пам'яті, що не є суміжними
• Можливість динамічного додавання або видалення вузлів із розгортань Kubernetes без порушення поточних операцій
• Підтримка шаблонів зв'язку «рівний з рівним» (P2P)

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

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

Сучасні високопродуктивні мережі

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

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

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

Філософія проєктування високопродуктивних мереж

Створення високопродуктивних мережевих систем вимагає переосмислення кількох фундаментальних припущень:

1. Володіння буфером: На відміну від традиційних сокетів, де ядро керує мережевими буферами та вимагає копіювання між простором користувача та простором ядра, RDMA вимагає, щоб програми керували власними буферами. Коли програма ініціює мережеву операцію, вона передає володіння буфером мережевій карті до завершення операції, що усуває потребу в копіюванні даних.
2. Реєстрація пам'яті: Програми повинні зареєструвати області пам'яті в ядрі операційної системи. Ядро налаштовує відображення віртуальних адрес, які дозволяють процесору, графічним процесорам і мережевим картам розуміти однакові віртуальні адреси. Ця реєстрація є одноразовою операцією, яка уможливлює подальшу передачу даних без копіювання.
3. Площина управління проти площини даних: Високопродуктивні мережі розділяють операції на дві категорії:
   1. Операції площини управління (як-от налаштування з'єднання та реєстрація пам'яті) проходять через ядро для забезпечення безпеки
   2. Операції площини даних (фактична передача даних) обходять ядро для максимальної продуктивності
4. Отримання перед передачею: Без буферів, якими керує ядро, програми повинні заздалегідь публікувати операції отримання, вказуючи, де мають бути розміщені вхідні дані. Це фундаментальна зміна в порівнянні з моделлю сокетів, де програми можуть отримувати дані в будь-який час.
5. Завершення на основі опитування: замість очікування сповіщень ядра за допомогою таких механізмів, як epoll, програми безпосередньо опитують апаратні черги завершення. Це усуває накладні витрати на системні виклики та дозволяє негайно реагувати на завершені операції.
6. Обізнаність про апаратну топологію: Розуміння та оптимізація апаратної топології мають вирішальне значення для досягнення максимальної продуктивності.

Розуміння апаратної топології

Екземпляри AWS p5 мають складну внутрішню архітектуру. Як показано нижче, кожен екземпляр містить два сокети CPU, що утворюють два вузли NUMA, причому кожен вузел NUMA з'єднаний з чотирма комутаторами PCIe:

Розгорнути

Під кожним комутатором PCIe ми знаходимо:

• Чотири мережеві карти EFA по 100 Гбіт/с
• Один GPU NVIDIA H100
• Один твердотільний накопичувач NVMe SSD

Шляхи передачі даних для передач TCP/IP і RDMA демонструють фундаментальні відмінності в їхніх підходах:

Розгорнути

З TCP/IP (ліва сторона) дані потрібно копіювати кілька разів:

1. З пам'яті GPU до буфера програми в просторі користувача
2. З буфера простору користувача в буфер ядра
3. З буфера ядра до мережевої карти
4. Зворотний процес відбувається на стороні приймача

Кожна операція копіювання споживає цикли CPU та пропускну здатність пам'яті. Додаток також повинен перемикати контекст між простором користувача та простором ядра для кожної мережевої операції.

На відміну від цього, RDMA (права сторона) забезпечує справжню передачу даних без копіювання (zero-copy):

1. Мережева карта зчитує дані безпосередньо з пам'яті GPU
2. Дані прямують безпосередньо до віддаленої мережевої карти
3. Віддалена мережева карта записує дані безпосередньо в цільову пам'ять GPU
4. Додатку потрібно лише перевірити чергу завершення (CQ) у просторі користувача, щоб дізнатися, коли передачу завершено

Завдяки RDMA та правильному апаратному з'єднанню передача даних між двома GPU вимагає проходження лише локального комутатора PCIe та мережі:

Розгорнути

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

Розгорнути

Створення за допомогою libfabric

Ми використовували libfabric — фреймворк, який надає загальний інтерфейс для служб fabric.Наша реалізація використовує два типи операцій RDMA:

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

Шлях до максимальної продуктивності

Наша розробка пройшла кілька етапів:

1. Реалізовано базову однонаправлену передачу повідомлень за допомогою SEND/RECV
2. Розширено до двостороннього зв'язку
3. Додано GPUDirect RDMA WRITE для прямої передачі між 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}
}