GPU-vezérelt kommunikációval gyorsítják az AI-betanítást a kutatók

A GPU-kommunikáció jelentős lassító tényező az AI-munkaterhelésekben. A kutatók adatai szerint a kommunikáció az AI-betanítás előrehaladásának 43,6%-át, a teljes folyamatnak pedig 32%-át emészti fel. Különösen a Mixture-of-Experts (MoE) modellek esetében ez az arány elérheti a 47%-ot.

UC Berkeley kutatói az UCCL projekt keretében bemutatták a mKernelt, egy új könyvtárat, amely persistent CUDA kerneljeivel egyesíti az intra-node NVLink kommunikációt, az inter-node RDMA-t és a számítást egyetlen kernelbe. Ezzel a GPU-vezérelt megközelítéssel a korábbi, CPU-vezérelt rendszerek korlátait hivatottak leküzdeni.

Kapcsolódó: gépi tanulás

A hagyományos, host-driven kommunikációban a CPU irányítja a folyamatot, ami a GPU-k számítási kapacitásának növekedésével nem skálázódik megfelelően. A GB300 NVL72 rack 72 Blackwell Ultra GPU-ja és 36 Grace CPU-ja mellett a CPU-vezérlés mikroszekundumos késleltetéseket okoz.

Kapcsolódó: LLM-gyorsítás

A számítás és kommunikáció szimfóniája

A mKernel GPU-vezérelt kommunikációt valósít meg, ahol a GPU indítja az átviteléket, és a kommunikáció ugyanabban a kernelben fut, mint a számítás. A könyvtár persistent CUDA kerneljei egyesítik az intra-node NVLink és az inter-node RDMA kommunikációt, lehetővé téve a számítás és a kommunikáció finom granularitású, tile/chunk szintű átfedését.

Kapcsolódó: PyTorch DDP

A fúziós kernelök hangja

A mKernel öt fő fúziós kernelt kínál: AllGather + GEMM, GEMM + AllReduce, MoE Dispatch + GEMM, Ring Attention, és GEMM + ReduceScatter. Ezek a kernelek speciálisan az AI-betanítás gyakori műveleteit optimalizálják, például a MoE tokenek útválasztását és a figyelmi mechanizmusok hatékonyabb végrehajtását. A kommunikációs backendet a libibverbs könyvtárra építették, függetlenül a NCCL vagy NVSHMEM használatától.

Kapcsolódó: diffúziós modellek

A rendszert két, egyenként 8 H200 GPU-val szerelt, 2-nodeos klaszteren tesztelték, eltérő inter-node hálózati konfigurációkkal. Az eredmények alapján a mKernel jelentős előrelépést kínál a GPU-kommunikáció hatékonyságában, különösen nagy léptékű AI-modellek betanítása során, mint például a GB300 NVL72 rendszeren, 2024-ben.

Kapcsolódó: CUDA-optimalizálás