Тиімді және портативті mixture-of-experts байланысы
Портативті Mixture-of-Experts (MoE) байланысының шолуы, ауқымды AI модельдеріндегі GPU параллелизмін оңтайландыруға және кідірісті азайтуға бағытталған
Біз стандартты All-to-All байланыс примитивтеріне қарағанда 10 есе жылдам өнімділікке қол жеткізетін Mixture-of-Experts (MoE) байланысына арналған жоғары өнімді, портативті, ашық бастапқы кодты кітапхананы ұсынамыз.
Біздің іске асыруымыз MoE байланысының жоғары тиімділігін қамтамасыз ететін бірнеше негізгі техникалық инновацияларды қамтиды:
GPU арқылы басталатын байланыс (IBGDA): CPU қатысуын айналып өту арқылы кідірісті айтарлықтай азайтып, тікелей GPU-дан NIC-ке байланысты қолдайды
Байланыс пен есептеуді қабаттастыру: Жеке жіберу және қабылдау кезеңдері бар бөлінген ядро архитектурасы есептеулердің желілік тасымалдаулармен қатар жүруіне мүмкіндік береді
Ең жылдам бір түйінді өнімділік: Бір түйінді конфигурацияларда бұрынғы ең жылдам іске асыруға қарағанда 2.5 есе төмен кідіріс
Тиімді және портативті көп түйінді өнімділік: Біздің іске асыруымыз стандартты all-to-all байланысына қарағанда 10 есеге дейін жылдамдыққа қол жеткізеді. Жоғары мамандандырылған іске асырулардан шамамен 2 есе баяу болғанымен, біздің тәсіліміз NVSHMEM версийлары мен желілік орталарда (NVLink, CX-7 және EFA) жақсырақ портативтілікті ұсынады
Кітапхана толығымен ашық бастапқы кодқа ие және https://github.com/ppl-ai/pplx-kernels мекенжайында қолжетімді.
Бұл мақалада біз ауқымды MoE модельдеріндегі сарапшылық параллелизм мәселелерін зерттейміз және таратылған GPU орталарында тиімді токендерді тарату және біріктіру тәсілімізді сипаттаймыз.
Кіріспе
DeepSeek R1 және Mixtral 8x7B сияқты Mixture-of-Experts (MoE) модельдері әрбір токен үшін белсендірілетін салмақ санын шектеу арқылы тығыз модельдерді жақсартады. Әрбір декодер қабатының көп қабатты перцептронындағы (MLP) қабаттар санын немесе сызықтық қабаттардың өлшемін ұлғайту арқылы жоғары параметр санына қол жеткізудің орнына, MoE модельдері дәстүрлі MLP-ті бірнеше сарапшы және мартируатормен (router) алмастырады. Мысалы, DeepSeek R1-дің 671B параметрінің ішінен шығару кезінде берілген токенмен тек 37B көбейтіледі. Бұл токенді декодтау үшін қажетті есептеу көлемін айтарлықтай азайтады және тығыз модельмен салыстырғанда кідірісті жақсартады.
MoE модельдері тығыз модельдермен салыстырғанда шығару үшін кейбір қосымша қиындықтар тудырады. Сарапшылардың өздері шағын MLP қабаттары болғанымен, әрбір декодер қабаты токеннің қандай сарапшыларға жіберілетінін анықтайтын маршрутизаторды қамтиды, ал әрбір токен бірнеше сарапшыға жіберіледі. Маршрутизатор әдетте ықтималдықтар бөлінісін шығаратын шағын сызықтық қабат болып табылады. Әдетте ең жоғары K ұпайлары бар сарапшылар таңдалады және соңғы белсендіру маршрутизатордың ықтималдықтар бөлінісінен алынған салмаққа көбейтілген сарапшы шығыстарын қосып, өлшенген орташа мән ретінде есептеледі.
Таратылған MoE жүйелеріндегі кідірісті азайту үшін параллелизмді бірнеше құрылғы арқылы пайдалануға болады, бірақ бұл байланыс қиындықтарын тудырады. Mixtral 8x7B немесе Llama-70B сияқты модельдер жылдам және төмен кідірісті NVLink өзара байланыстарының (900 Гб/с дейін) артықшылығын пайдалана отырып, бір түйіндегі 8 құрылғыға сыяды. Дегенмен, үлкенірек модельдер InfiniBand (400 Гб/с шегінде) көмегімен сарапшыларды бірнеше түйінге бөлуді талап етеді, бұл қосымша кідіріс қиындықтарын тудырады.
Бұл мақалада біз сарапшылық параллелизм мәселесін зерттейміз, тағайындалған сарапшыларды ұстайтын құрылғыларға токендерді тиімді жіберетін және сарапшы шығыстарының өлшенген қосындысын есептеу арқылы оларды қайта жинайтын GPU ядроларын іске асыруымызды сипаттаймыз. Белгілі бір желілік жабдыққа бағытталған агрессивті түрде оңтайландырылған іске асырулар бар болғанымен, біз NVSHMEM примитивтерінің минималды жиынтығына сүйене отырып, бірақ стандартты all-to-all байланысына қарағанда 10 есе жылдам өнімділікті қамтамасыз ететін жоғары портативті іске асыруды ұсынамыз. Біздің іске асыру әр түрлі желілік орталарда тамаша портативтілікті сақтай отырып, бір түйінді конфигурацияларда заманауи өнімділікке қол жеткізеді.
Бөліктерге бөлу және параллелизм
Тиімді шығару үшін модельдің салмақтары құрылғы жадында сақталуы тиіс, сонымен бірге модель арқылы тура өту үшін қажетті белсендірулерге, KV кэштеріне және басқа буферлерге жеткілікті кеңістік қалдырылуы керек. Ең қабілетті модельдер ең қабілетті GPU-лардың да сыйымдылығынан асып түседі, сондықтан шығару модельдің салмақтарын бірлесіп сақтай алатын бірнеше құрылғыға таратылуы тиіс. Салмақтарды бөлу схемаларына сүйене отырып, есептеулерді түсіру және құрылғыларды синхрондау үшін әртүрлі байланыс және есептеу схемалары қолданылуы тиіс.
Бірінші суретте көрсетілген Сарапшылық параллелизм (EP) тек сарапшылық есептеулерді параллельді етеді. Әр түрлі сарапшылар өздерінің салмақтарын ұстайтын әртүрлі құрылғыларға тағайындалады. Бағыттаудан кейін токендер сәйкес құрылғыға жіберіледі, ал нәтижелер кейіннен жиналып, жинақталады. Бағыттаудың күрделілігі модельдің басқа бөліктеріндегі параллелизм дәрежесіне байланысты болады: басқа қабаттарды қайталау бағыттау қажеттілігін толығымен жоюы мүмкін, өйткені әрбір ранг жергілікті қайталанатын бағыттау кестесінен токендерді таңдай алады. Дегенмен, егер рангтердің тек біреуі ғана бағыттауды орындаса, токендерді, индекстерді және салмақтарды өздерінің тиісті сарапшыларына тарату үшін трансляция қажет болады. Соңында, all-gather немесе all-to-all трансляциясы шығыс токендерін модельдің орындалуын жалғастыратын кез келген рангпен синхрондайды. Мұндай іске асыру салыстырмалы түрде қарапайым, өйткені torch қажетті примитивтерді ұсынады, бірақ жинақтауды сарапшылардан токендерді жинайтын жинау операциясымен біріктіру арқылы кейбір артықшылықтарға қол жеткізуге болады.
Кластердегі түйіндер модель сарапшылық емес қабаттарды (мысалы, назар аудару, норма және іріктеу) орындап жатқанда бос тұруы мүмкін болғандықтан, тек сарапшылық параллелизм мінсіз түрде ауқымданбайды. Дегенмен, бұл қабаттардың есептеулері, ең алдымен назар аудару, Tensor Parallelism (TP) артықшылықтарын пайдалана алады. Көптеген модельдер көп басты назар аударуға сүйенеді, бұл назар аудару бастары мен оларға сәйкес Q, K және V проекциялары әртүрлі қабаттар арасындағы бөліктерді қайталау немесе жинау арқылы құрылғылар арасында бөлінуі мүмкін екенін білдіреді. Егер назар аудару барлық құрылғыларға таратылса, all-gather белсендірулерді синхрондай алады, бұл тиімді іске асыру үшін сарапшылық параллелизм жағдайына ұқсас синхрондау примитивтерін талап етіп, бағыттауды қайталауға мүмкіндік береді. Дегенмен, бұл деңгейдегі параллелизмге шектеулер бар, өйткені назар аудару бастарының санын белгілі бір шектен төмен азайту қайтарымның кемуіне әкеледі.
Барлық құрылғыларды барынша пайдалану және 128 немесе 256 GPU-ға дейінгі сарапшылық параллелизмнің өте жоғары дәрежесін қолдау үшін Деректер параллелизмі (DP) қажет. Бұл схема бойынша модель бөлінген құрылғылар сұраныстарды бір мезгілде өңдеу, назар аударуды есептеу және олардың жергілікті тобында бөлінген KV кэштерін сақтау үшін топтастырылады. Осы параллель топтардың бірнеше данасы әрқайсысы сарапшылық салмақтардың әртүрлі ішкі жиынтын орналастыра отырып, сарапшыларды бағалауға бірлесіп қатысады. Назар аудару бастарының санына байланысты топ әдетте бүкіл түйіннің өлшеміне дейін созылуы мүмкін, өйткені түйінішілік байланыс жылдамырақ болады. Мысалы, жоғарыдағы суретте бір DP рангі екі сұранысқа (A және B) тәуелсіз қызмет көрсетеді, бұл сұраныстар үшін назар аударуды, норманы және кез келген басқа есеп жүргізуді өңдейді. Басқа DP рангі бөлек C сұранысын өңдейді. Дегенмен, бірінші түйін сарапшылардың жартысын орналастырады, ал екінші түйін қалған жартысын орналастырады, сондықтан A және B токендерін бағыттағаннан кейін екінші түйінге жіберілуі мүмкін және керісінше. Бұл сирек байланыс мәселесіне әкеледі: әрбір құрылғы кез келген басқа мақсатты рангке әртүрлі сандағы токендерді жібере алады. Torh-тың бар примитивтері, ең алдымен all_to_all, онша жақсы бейімделмеген, өйткені олар толтырудың (padding) немесе GPU-CPU синхрондауының және метадеректерді трансляциялаудың қандай да бір түрін талап етуі мүмкін. Байланысты тиімді іске асыру үшін CPU араласуынсыз бағыттау ақпаратын өңдеу және құрылғылардан байланысты бастау үшін реттелетін ядролар қажет. Бағыттаудан кейін диспетчерлік ядро токендерді олар бағытталған рангтерге жіберуі тиіс, ал біріктіру жағында ағымдағы DP тобындағы сұраныстарға тиесілі белсендірулер жиналуы керек. Сонымен қатар, трансляцияны азайту және кейінгі қабаттар үшін барлық рангтердегі токендерді дұрыс қайталау үшін бірнеше құрылғысы болуы мүмкін DP тобында жұмыс теңгерімделуі тиіс.
NVSHMEM
NVSHMEM — бұл NVIDIA-ға тән OpenSHMEM іске асыруы, ол негізіндегі аппараттық құралдың күрделілігін абстракциялайтын, құрылғылар арасындағы тасымалданатын байланыс мүмкіндіктерін ұсынады. API құрылғыдан құрылғыға оқу және жазу әрекеттерін сипаттай алады, олар жеке тасымалдау қабаттарының примитивтеріне сәйкестендіріледі. NVLink те, RDMA да қолдау табады, бұл белгілі бір деңгейде тасымалданғыштық беріп, NVSHMEM ядроларына бір түйін ішінде немесе бірнеше түйіні бар кластер бойымен жұмыс істеуге мүмкіндік береді. Біздің ядроларда біз деректерді тасымалдау және қатысушы құрылғыларды синхрондау үшін қашықтағы жадыға жазу операциялары мен сигнал беру операцияларын қолданамыз.
NVSHMEM операциялары симметриялық жад тұжырымдамасының төңірегінде құрылған: олар құрылғы аралық байланысқа қатысатын барлық құрылғыларда бөлінген буферлерде жұмыс істейді. Жергілікті құрылғы операцияның мақсатын көрсету үшін нысана рангтің индексімен қатар өз алдынғы жергілікті буферінде алынған мекенжайды пайдалана алады. Төмендегі сурет бұл тұжырымдаманы көрсетеді: екі GPU де өлшемі бірдей симметриялық буферлерді бөледі, оларда src көрсеткіштерін сақтайды. Бірінші ранг екінші рангке 3 бүтін сан жібергісі келеді, оларды буфердің басына орналастырады: nvshmem_int_put_nbi нысана құрылғыны көрсете отырып, бастапқы мекенжайды жергілікті буферден алады. Екінші ранг өз буферінен офсетті алады, бірінші құрылғыға бір элементті жібереді, бір элементке ығыстырады. Нысана мекенжайлар әрқашан nvshmem_alloc көмегімен бөлінген симметриялық буферлер болуы тиіс болса, бастапқы буферлер NVSHMEM-де алдын ала тіркелген жағдайда құрылғы жадының еркін аймақтары бола алады.
NVSHMEM примитивтердің кең ауқымын қамтамасыз етсе де, біздің ядроларымыз барлық синхрондау мен қоршауды оларға негіздей отырып, тек 3 функцияға сүйенеді.
nvshmemx_putmem_signal_nbi_warp: Деректер блогын бір құрылғыдан екіншісіне тасымалдайды, сонымен бірге қашықтағы құрылғыда жалаушаны орнатады. Операция 64 биттік орынды орнатады (NVSHMEM_SIGNAL_SET) немесе ұлғайтады (NVSHMEM_SIGNAL_ADD). Жалауша жадтың бүкіл блогы тасымалданғаннан кейін жаңартылады. Егер қашықтағы құрылғы жалаушаның өзгерісін байқаса, ол өзінің жадындағы буферге қауіпсіз түрде қол жеткізе алады. Бұл функция деректерді тасымалдауды және синхрондауды біріктіру үшін пайдалы.
nvshmemx_signal_op бір жад орнында, әдетте 64 биттік жалаушада жұмыс істейді, оны автоматты түрде орнатады немесе ұлғайтады. Бұл метадеректерді жіберуде және құрылғыларды синхрондауда пайдалы.
nvshmem_uint64_wait_until қашықтағы құрылғы оны жаңартқанша жалаушаны сұрау үшін сигналды қабылдау жағында қолданылады.
nvshmem_quiet немесе nvshmem_fence арқылы операциялар арасындағы реттіліктің кез келген түрін қамтамасыз ету қымбатқа түседі, сондықтан біздің ядроларды іске асыруымыз кез келген басқа NVSHMEM примитивтерін пайдаланудан аулақ болады және жеке жіберу-қабылдау жұптарынан басқа кез келген тәуелділіктерді азайтады. Деректерді жіберген кезде біз әрқашан деректердің жергілікті рангтен тіпті жіберілуін күтпестен, функциялардың блокталмайтын нұсқасын пайдаланамыз. Ядролар тобы бойынша басқа, имплицитті синхрондау NIC-терге оларға қол жеткізу қажет болған кезде бұл буферлердің бұзушы түрде қайта жазылмауын қамтамасыз етеді.
NVSHMEM GPUDirect RDMA қолданады, осылайша NIC CPU-ны қатыстырмай деректерді алмасу үшін GPU жадынан тікелей оқи және жаза алады. Сонымен қатар, ConnectX NIC-терінде NVSHMEM GPU арқылы басталатын байланысты (Infiniband GPUDirect Async немесе IBGDA ретінде де белгілі) қолдайды, бұл GPU-ға RDMA операцияларын тікелей NIC-ке жіберуге және NIC-тің аяқтау кездерін тікелей сұрауға мүмкіндік береді. GPU арқылы басталатын байланысты қолдамайтын платформаларда NVSHMEM хост CPU-да байланысты бастау және GPU атынан аяқталуды сұрау үшін "прокси" тредин жасақтайды. GPU арқылы басталатын байланыстың қолданылатынына немесе қолданылмайтынына қарамастан NVSHMEM бағдарламасы портативті болып табылады. Дегенмен, GPU арқылы басталатын байланыс CPU-ға айналып өтуді толығымен жоятындықтан кідірісті айтарлықтай қысқартады.
Портативті ядролар
Біз MoE байланысын бір жұп диспетчерлік және біріктіру ядролары арқылы жүзеге асырамыз. Диспетчерлік ядролар әр рангтегі токендер мен бағыттау ақпаратын оқуға, оларды тиісті сарапшыларға жіберуге жауапты. Біріктіру ядролары сарапшылар шығарған белсендірулерді жинайды және оларды бастапқы рангтеріне кері жібереді, сонымен бірге маршрутизатор есептеген салмақтар негізінде таңдалған сарапшылардан өлшенген орташа мәнді есептейді. Ядролар деректерді тасымалдауларға есептеулермен қабаттасуға мүмкіндік беру үшін жіберу және қабылдау компонентіне одан әрі бөлінеді. Жіберу ядролары толығымен блокталмайтын және синхрондалмайтындар: олар тек барлық жазуларды қашықтағыларға жібереді. Басқа жағынан, біріктіру ядролары барлық қажетті деректер тасымалданғанша күтіп, тек жадтан оқиды. Жұмысты NIC-терге жібергеннен кейін, деректер сым арқылы асинхронды түрде тасымалданып жатқанда, GPU жергілікті деңгейде басқа пайдалы жұмыстарды орындай алады, мысалы, ортақ сарапшыларды қолдану немесе назар аударуды есептеу.
Әрбір ядролар жұбының өзінің симметриялық жад буферлері бар, олардың әрбір рангі барлық басқа рангтерден деректерді алу үшін жеке сақтау орнын бөледі. DP тобынан жіберілуі мүмкін токендердің максималды саны рангтер арасында тіркелген, бұл сонымен қатар әрбір ранг әрбір DP рангінен әрбір жергілікті сарапшы үшін ала алатын жоғарғы шекті белгілейді. Бұл жіберуші рангтерге олардың арасында ешқандай синхрондауды талап етпестен, жазу үшін нысана рангте бірегей мекенжайды алуға мүмкіндік береді. Деректер қабылданғаннан кейін ядролар қабылдаушы ядролар талап ететін оңтайлы форматта жадта орналастыру үшін токендерді ауыстырады. Буферлердің өлшемі үлкен болғанымен, олар модельдің барлық тізбекті қабаттарында қайта пайдаланылады.
Жоғарыдағы суретте көрсетілгендей, жаһандық синхрондаудың жалғыз түрі біріктіру-қабылдау ядросында іске асырылады. Деректер жіберуден жіберілгеннен кейін, ранг басқа рангтерден қажетті барлық токендерді алған кезде оралады және орындалуын жалғастырады және олар өңделген бойда біріктіру жіберу ядроларына кіреді. Біріктіру-қабылдаудағы тосқауыл ешбір рангтің басқа ранг әлі де деректерді алуды күтіп тұрған кезде алға кетіп, диспетчерлік жіберуді бастауына жол бермейді, бұл ықтимал бұзушы қабаттасуға әкелуі мүмкін. Синхрондау қарапайымдылық үшін біріктіру ядросында жасалады, оңтайландыру ретінде дәл сол тосқауыл диспетчерлік қабылдау және біріктіру жіберу ядролары арасында іске асырылуы мүмкін, бұл GPU басқа ядроларды іске қосып жатқанда NIC-терге тасымалдауларды реттеуге көбірек уақыт беру арқылы кідірісті одан әрі азайтады.
Екі ядро да құрылғылардың барлық қолжетімді SM-деріне бөлінген: диспетчерлік жіберу және біріктіру қабылдау ядролары бір рангке шаққандагы максималды токен саны (max_m) бойынша параллелизациялануы тиіс болса, диспетчерлік қабылдау және біріктіру жіберу ядролары әр ранг пен жергілікті сарапшыдан сонша токенді өңдеуі тиіс, ең көбі max_m * (num_experts // EP) * (EP // TP).
Диспетчерлеу
Диспетчерлік ядро токендер бағытталған сарапшыларды орналастыратын пирлерге токендерді және блок масштабтау факторларын жіберуге жауапты. Жіберуші жағы екі тапсырманы параллельді ету үшін warp мамандануына сүйенеді: әр рангтен қанша токен жіберілетінін анықтау үшін бағыттау ақпаратын біріктіру және токендерді қашықтағыларға жіберілетін хабарларға буып-түю. Қабылдаушы жағы алдымен барлық токен сандарының қабылдануын күтеді, содан кейін ол ортақ жадтан жүктемелерді оқып, төменгі ағындағы топтастырылған GEMM ядроларына берілетін тензорларға ағытады. Ол сондай-ақ әрбір токеннің қайда қайтарылуы керектігін көрсету үшін біріктіру ядроларымен ортақ буферлерде ақпаратты сақтайды. Бұл ақпарат қабылдау ядросы токендерді анықталмаған тәртіпте іргелес буферлерде ауыстырған кезде қажет болады.
Жіберуші бөлігінде координирленген әрбір блоктағы 32 тредтен тұратын жалғыз warp бағыттау ақпаратын оқуға және әрбір мақсатты сарапшы осы рангтен алуы тиіс токендер санын санауға, nvshmem_signal_op көмегімен сан мен бірді жіберуге жауапты. Саны бірге ұлғайтылады, өйткені қашықтағы жақта нөлден нөл емеске ауысу бастапқы рангтен сандардың алынғанын білдіреді. Параллель түрде қалған warp-тар токендерді барлық блоктарда симметриялық жадқа параллель түрде көшіру үшін бірлесіп жұмыс істейді, белсендірулерді, масштабтау факторларын және жергілікті рангтегі индексін іргелес бөлікке буып-түеді. Индекс біріктіру жіберушісіне токен жазылатын мекенжайды анықтау үшін қажет. Әрі қарай, warp топтары арасындағы тосқауыл арқылы барлық деректердің көшірілгеніне көз жеткізгеннен кейін, warp-тар әлі де тәуелсіз жұмыс істейді, олардың әрқайсысы бірдей буферді параллель түрде әр түрлі сарапшыға жібереді. Токендер сонымен қатар қашықтағы құрылғыда жергілікті рангтен жіберілген токендер санын автоматты түрде арттыратын nvshmemx_putmem_signal_nbi_warp көмегімен жіберіледі. DP тобы ішінде, әрбір ранг жіберілетін токеннің көшірмесіне ие болғандықтан, диспетчерлік біркелкі теңгерімделген, әрбір құрылғы токендердің ішкі жиынтын жібереді.
Қабылдау жағында барлық ядролар алдымен жалпы токен сандарының жіберілуін күту арқылы жіберуші рангтермен синхрондайды. Одан кейін олардың барлығы putmem шақыруының семантикасының арқасында бастапқы рангтерден келген барлық жүктемелердің де жіберілгенін көрсететін автоматты түрде ұлғайтылған жіберілген токен сандарының жалпы сандарға дейін орнығуын күтеді. Ядролар операцияны барлық блоктар мен тредтер бойынша параллелизациялай отырып, nvshmem_uint64_wait_until арқылы сандарды сұрайды. Одан кейін блокаралық тосқауыл барлық деректер дұрыс қабылданбаса, ешбір блоктың буферлерді оқымауын қамтамасыз етеді. Блоктарға бөлінген және автоматты түрде ұлғайтылған есептегіш арқылы синхрондалған токендер симметриялық жадтан сарапшыларды іске асыратын төменгі ағынды ядролардың кірістеріне айналатын тұрақты іргелес буферлерге көшіріледі. Бастапқы ранг, сарапшы индексі және токен индексі бөлек сақталады, бұл біріктіру ядросы белсендірулерді жіберуі тиіс орынды дәл көрсетеді. DP тобы ішіндегі токендер әртүрлі құрылғылардан жіберілсе де, олардың барлығы сәйкес сарапшыға жіберілу үшін бірге топтастырылады.
Біріктіру
Біріктіру ядролары қайтадан жіберу және қабылдау жартыларына бөлінеді: жіберушілер квантталмаған 16 биттік белсендірулерді қашықтағы құрылғыларға көшіреді, ал қабылдау ядролары деректердің жіберілуін күтеді және барлық сарапшы үлестерінің өлшенген орташа мәнін жергілікті түрде есептейді. Сонымен қатар, олар сондай-ақ диспетчерлік-біріктіру тізбегін синхрондау үшін тосқауыл ретінде әрекет етеді: әрбір ранг жіберу ядроларына кірген кезде әр пирде жалаушаны орнатады, ал қабылдау ядроларына жалаушаның орнатылғанын байқамаса, қайтаруға рұқсат етілмейді. Синхрондау кідірісі минималды, өйткені ол нақты байланыспен және есептеумен қабаттасады.
Жіберуші жағында ядролар барлық жергілікті сарапшыларға тағайындалған токендер тізімін параллель түрде аралап, оларды нысана рангтегі буферге жазады. Нысана ранг, сарапшы индексі және токен индексі токендер алынғаннан кейін скатер ядролары толтырған токен бойынша буферлерден оқылады. Жоғарыдағы суретте көрсетілгендей, әрбір жіберушінің жазу үшін әр сарапшы үшін жеке жад аймағы бар, бұл синхрондау қажеттілігінен аулақ болады. Диспетчерлеуге ұқсас, біріктіру деректердің алынғанын көрсету үшін нысана рангтегі токен есептегіштерін автоматты түрде ұлғайтады: есептегіш токен жіберілген сарапшылар санына сәйкес келгенде, токен мазмұнына қол жеткізуге болады.
Қабылдау ядросында токендер тізімі сигналы операциясы орнатқан жалаушаны сұрау арқылы олардың мазмұнының келуін күтіп, бірнеше блоктар бойынша параллель түрде араланады. Келген кезде жүктемелер жеке буферлерден оқылады, ал бағыттау кестесі қай буферден оқу керектігін және әр сарапшыға қандай салмақ тағайындау керектігін көрсетеді. Нәтижелер сырттан бөлінген тензорларға жазылады, ядро барлық құрылғылар тосқауылдан өткеннен кейін орындалуын аяқтайды. DP тобы бойынша барлық рангтер өздерінің көшірмелерін есептеу үшін әрбір сарапшы белсендіруінің көшірмесін алады.
Бенчмарктар
Біз ядроларымызды InfiniBand және CX-7 NIC арқылы қосылған 16 түйінге таратылған 128 GPU кластерінде бағалаймыз. Біз GPU-initiated Async (IBGDA) пен CPU проксиі бар сенімді қосылымды (RC) екеуін де сынап көреміз. Біз оларды жоғары оңтайландырылған DeepSeek іске асыруымен, сондай-ақ PyTorch (NCCL арқылы) немесе NVSHMEM ұсынатын тығыз примитивтермен салыстырамыз.
Тығыз NVSHMEM all-to-all операциясында IBRC және IBGDA өнімділігі ұқсас болғанымен (6378 µs vs 6180 µs), IBGDA сирек ядролармен айтарлықтай жылдамырақ. Диспетчерлеу мен біріктіруді қосқанда, IBGDA 902 µs пайдаланады, ал IBRC 3223 µs алады — кідірісте 3.6 есе жақсарту. All-to-all өткізу қабілетіне шектелген болса, сирек ядролар кідіріспен шектеліп, әлдеқайда аз деректерді трансляциялайды. GPU пен NIC-ті үйлестіру үшін CPU проксиін талап етпей, желілік тасымалдауды тікелей GPU-дан іске қосу арқылы ұшты-ұштық кідіріс айтарлықтай азайтылады.
Біздің портативті ядроларымыз жоғары оңтайландырылған DeepSeek ядроларынан шамамен 2 есе баяу болғанымен, олар тығыз ядролармен салыстырғанда кідірісті 10 есе жақсартады (902 µs vs 9944 µs).
Сонымен қатар, жіберуші-қабылдаушы компоненттеріне бөлу кітапхана примитивтерінен айырмашылығы, кідірістің бір бөлігін жасыруға мүмкіндік береді.
Бір түйінді (EP8) жағдайда NVSHMEM тасымалдау үшін NVLINK пайдаланады, бұл түйінаралық желіге қарағанда төмен кідіріс пен жоғары өткізу қабілетін қамтамасыз етеді. Біздің портативті ядроларымыз бір түйінде DeepEP-ке қарағанда шамамен 2.5 есе жылдамырақ (186 µs vs 481 µs).
Қосымша жақсартулар
Мұнда сипатталған ядролар ML фреймворктерінің кірістірілген примитивтерінен асып түседі және NVLink немесе InfiniBand сияқты арнайы құрылғы аралық тасымалдаулар үшін тым маманданбай-ақ лайықты өнімділікті ұсынады. Бұрын айтылған мүмкіндіктерден басқа, байланыс примитивтерін неғұрлым мамандандырылған нұсқалармен ауыстыру арқылы қосымша өнімділік ұтымдылығына қол жеткізуге болады. Мысалы, NVLink бойынша симметриялық жадты пайдалануды құрылғылар арасында ортақ буферлермен алмастыруға болады, бұл кейбір көшіруді жояды және токен топтарының орнына жеке токендер арқылы неғұрлым дәлірек синхрондау схемаларына мүмкіндік береді. InfiniBand бойынша іске асыру түпкілікті кезек жұптарына тікелей қол жеткізе алады, бұл күтілетін және жіберілген токен сандарын тасымалдау қажеттілігін жояды және токендер келген кезде оларды кезектерден шығыс буферлеріне тазартады. Дегенмен, мұндай іске асырулар портативтілік құнымен келеді, ал біздің іске асыру белгілі бір аппараттық құралдар үшін оңтайландыруға әзірлеу күш-жігерін жұмсамастан бұрын нақты конфигурациялары бар кластерлерді бағалауға мүмкіндік береді.
Қорытынды
Біз әртүрлі аппараттық конфигурациялар бойынша үйлесімділікті сақтай отырып, стандартты all-to-all байланысымен салыстырғанда 10 есе жылдам өнімділікке қол жеткізетін MoE байланысына арналған жоғары өнімді, портативті кітапхананы ұсындық. NVLink бар бір түйінді орналастыруларда біздің шешіміміз бұрынғы іске асыруларға қарағанда 2.5 есе төмен кідірісті көрсетеді.
Біздің тәсіліміз GPU арқылы басталатын байланыс қолдауы, есептеу-байланыс қабаттасуын қамтамасыз ететін бөлінген ядро архитектурасы және минималды NVSHMEM примитивтерін пайдаланатын тиімді токен диспетчерлігі сияқты негізгі инновациялар арқылы өнімділікті портативтілікпен теңестіреді. Көп түйінді орнатуларда жоғары мамандандырылған іске асырулардан шамамен 2 есе баяу болғанымен, біздің кітапханамыз әртүрлі желілік орталарда (NVLink, CX-7 және EFA) жоғары икемділікті ұсынады.
MoE модельдері ауқымдануын жалғастырған сайын, біздікі сияқты тиімді байланыс стратегиялары практикалық орналастыру үшін маңызды бола түседі. Біздің толығымен ашық бастапқы кодты іске асыруымыз https://github.com/ppl-ai/pplx-kernels мекенжайында қолжетімді.
Әдебиеттер тізімі
Дәйексөз
@online{perplexity2025EfficientAndPortableMixtureOfExpertsCommunication,
author = {{Perplexity} and Nandor Licker and Lequn Chen and Phi Hoang},
title = {{Efficient and portable mixture-of-experts communication}},
date = {2025-04-02},
url = {https://www.perplexity.ai/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}