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

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

Техническото предизвикателство

Нашият случай на употреба постави няколко ключови технически изисквания:

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

Въпреки че библиотеката 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 позволява директен трансфер на данни между паметта на графичния процесор и мрежовата карта, като напълно заобикаля CPU.

Философия на проектирането на мрежи с висока производителност

Изграждането на мрежови системи с висока производителност изисква преосмисляне на няколко фундаментални предположения:

1. Собственост на буфера: За разлика от традиционните сокети, където ядрото управлява мрежовите буфери и изисква копиране между потребителското пространство и пространството на ядрото, RDMA изисква приложенията сами да управляват своите буфери. Когато приложението инициира мрежова операция, то прехвърля собствеността върху буфера на мрежовата карта, докато операцията приключи, което елиминира необходимостта от копиране на данни.
2. Регистрация на паметта: Приложенията трябва да регистрират регионите на паметта в ядрото на операционната система. Ядрото настройва картографирания на виртуални адреси, които позволяват на CPU, графичните процесори и мрежовите карти да разпознават едни и същи виртуални адреси. Тази регистрация е еднократна операция, която дава възможност за последващи трансфери на данни без копиране.
3. Контролен план спрямо план за данни: Мрежите с висока производителност разделят операциите на две категории:
   1. Операциите на контролния план (като настройка на връзката и регистрация на паметта) преминават през ядрото, за да се гарантира сигурност
   2. Операциите на плоскостта за данни (действителният трансфер на данни) заобикалят ядрото за максимална производителност
4. Приемане преди предаване: Без буфери, управлявани от ядрото, приложенията трябва предварително да публикуват операции за приемане, като указват къде да бъдат поставени входящите данни. Това е фундаментална промяна спрямо сокетния модел, при който приложенията могат да получават данни по всяко време.
5. Завършване въз основа на заявки (Poll-based Completion): Вместо да чакат известия от ядрото чрез механизми като epoll, приложенията директно заявят хардуерните опашки за завършване. Това елиминира овърхеда от системните извиквания и позволява незабавна реакция на завършени операции.
6. Осъзнаване на хардуерната топология: Разбирането и оптимизирането за хардуерната топология е от решаващо значение за постигане на максимална производителност.

Разбиране на хардуерната топология

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

Уголеми

Под всеки PCIe превключвател намираме:

• Четири 100 Gbps EFA мрежови карти
• Един NVIDIA H100 GPU
• Един NVMe SSD

Пътищата на данните за TCP/IP и RDMA трансферите демонстрират фундаментални разлики в техните подходи:

Уголеми

При TCP/IP (лява страна) данните трябва да се копират многократно:

1. От паметта на GPU до буфера на приложението в потребителското пространство
2. От буфера в потребителското пространство до буфера на ядрото
3. От буфера на ядрото до мрежовата карта
4. Обратният процес се случва от страната на получаване

Всяка операция по копиране консумира процесорни цикли и честотна лента на паметта. Приложението също така трябва да превключва контекста между потребителското пространство и пространството на ядрото за всяка мрежова операция.

За разлика от тях, RDMA (дясна страна) осигурява истински трансфер на данни без копиране (zero-copy):

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

С RDMA и правилно хардуерно сдвояване, прехвърлянето на данни между два графични процесора изисква само преминаване през локалния PCIe превключвател и мрежата:

Уголеми

За разлика от тях, TCP/IP трансферите трябва да копират данни многократно през основната памет, което води до значително задръстване на PCIe шината:

Уголеми

Изграждане с libfabric

Използвахме libfabric - рамка, която предоставя общ интерфейс за мрежови услуги (fabric services). Нашата реализация използва два типа 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. Многопонишност: Специализирани теми за мрежовите операции на всеки графичен процесор.
4. Закачане към процесорни ядра (CPU Core Pinning): Обвързване на темите към конкретни CPU ядра, за да се избегнат NUMA ефектите и пропуските в кеша.
5. Шардиране на състоянието: Намаляване на споровете между темите.
6. Пакетно обработване на операциите: Подаване на множество операции заедно.
7. Мързеливо пускане на операциите: Операциите първо се поставят на опашка в приложението. След заявка (polling) на опашката за завършване, се опитваме да подадем чакащите операции към мрежовата карта, осигурявайки ефективно използване на мрежовите ресурси.
8. Разпределение на ресурсите с отчитане на NUMA: Разпределяне на libfabric ресурси като опашки за завършване на правилния NUMA възел, за да се минимизира латентността на достъпа до паметта.

Чрез тези оптимизации постигнахме крайна производителност от 3108 Gbps - 97,1% от теоретичния максимум на честотната лента.

Видеото по-долу показва нашата програма за команден ред в действие. Тя прехвърля несвързани части от GPU памет към отдалечен възел, постигайки скорост на предаване от 3108,283 Gbps - демонстрирайки близко до теоретичното използване на честотната лента на мрежовата инфраструктура:

Заключение

Изграждането на мрежова система с висока производителност изисква разбиране както на хардуерната архитектура, така и на принципите на системния дизайн. Въпреки че библиотеки като 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}
}