Болжау арқылы Sonar модельдерін жылдамдату

Спекулятивті кодтан шығару үлкен тілдік модельдердің (LLM) генерация жылдамдығын жылдам және шағын бастапқы модельдің көмегімен арттырады, ол ірірек мақсатты модель арқылы тексерілетін аяқтау кандидаттарын шығарады. Бұл әдіс бойынша қымбат мақсатты модельдің бір токенді шығаруының орнына, бір қадамда бірнеше токен шығарылады. Мұнда біз Perplexity жүйесінде Sonar модельдеріндегі токендер арасындағы кідірісті азайту үшін қолданылатын спекулятивті кодтан шығарудың әртүрлі түрлерінің іске асыру мәліметтерін ұсынамыз.

Спекулятивті кодтан шығару

Спекулятивті кодтан шығару токендер генерациясын жылдамдату үшін табиғи тілдердің құрылымын және трансформерлердің авторегрессиялық сипатын пайдаланады. Llama-70B сияқты ірі модельдер Llama-1B сияқты шағын модельдерге қарағанда көбірек білімге ие болғанымен, кейбір қарапайым тапсырмаларда олардың жұмысы ұқсас болады. Бұл ұқсастық кейбір тізбектерді аз шығынды модельдер жақсырақ генерациялайтынын, ал күрделі мәселелерді ірі модельдерге қалдыратынын көрсетеді. Мәселе қай аяқтаулардың жақсырақ екенін және шағын модельдің генерация сапасы ірі модельдікімен бірдей екенін анықтауда жатыр.

Бақытымызға орай, LLM — бұл авторегрессиялық трансформерлер: оларға токендер тізбегі берілгенде, олар келесі токеннің ықтималдық бөлінуін шығарады. Сонымен қатар, кіріс тізбегіндегі токендермен байланысты аралық мүмкіндіктерден алынған логиттер модельдің дәл сол токендерді шығару ықтималдығын да көрсетеді. Бұл қасиет болжам жасауға мүмкіндік береді: егер токендер тізбегі кіріс префиксінен бастап шағын модель арқылы генерацияланса, оның мақсатты модельмен қаншалықты сәйкес келетінін анықтау үшін оны ірі модель арқылы өткізуге болады. Кандидаттардың әрбір префиксі ықтималдықпен бағаланады және қабылдау шегінен асатын ең ұзын префикс таңдалады. Қосымша артықшылық ретінде мақсатты модель кейінгі токенді тегін береді: егер бастапқы модель n токен генерацияласа, бір қадамда n + 1 токенге дейін шығарылуы мүмкін.

Ұлғайту

Инференс кезінде спекулятивті іріктеу процесін шамамен 4 кезеңге бөлуге болады:

Префилл: KV кэш жазбаларын толтыру үшін мақсатты модель де, бастапқы модель де кіріс тізбегінде орындалуы тиіс. Medusa сияқты кейбір схемалар болжау үшін қарапайым тығыз қабаттарды пайдаланса, бұл мақалада біз өздерінің KV кэштерін қажет ететін трансформерге негізделген бастапқы модельдерге назар аударамыз.

Бастапқы генерация: бастапқы модель белгілі бір санның тіркелген токендерін шығару үшін қайталанады. Бастапқы тізбек сызықтық болуы мүмкін немесе модель берілген тереңдікке дейін ағаш тәрізді құрылысты зерттей алады (EAGLE, Medusa). Мұнда біз сызықтық тізбектерге назар аударамыз.

Қабылдау: мақсатты модель әрбір бастапқы токенге сәйкес келетін логиттерді құра отырып, бастапқы тізбекте жұмыс істейді. Ең ұзын қабылданатын тізбектің ұзындығы анықталады.

Мақсатты генерация: мақсатты модель логиттерді генерациялағандықтан, сәйкес келмейтін позицияда немесе тізбектің соңында логиттер кейінгі токенге сәйкес келеді. Бұл логиттер тізбекті қорытындылай отырып, мақсатты модельден сенімді токенді қамтамасыз ету үшін іріктелуі мүмкін.

Спекулятивті кодтан шығаруды іске асырудың әртүрлі әдістері бар. Бұл мақалада біз ішкі 1B моделі арқылы Sonar модельдерін жылдамдату үшін қолданған схемаларға, сондай-ақ DeepSeek ауқымындағы модельдерді жылдамдату үшін құрып жатқан болжау механизмдеріне назар аударамыз.

Мақсатты және бастапқы модель (Target-Draft)

Спекулятивті кодтан шығару кандидаттар тізбегін генерациялау үшін мақсатты модельге бастапқы модель ретінде қолданыстағы шағын LLM моделін қосу арқылы жүзеге асырылады. Өндірісте біз мақсатпен бірдей деректер жиынтығында баптаудан өткен Llama-1B моделі арқылы Sonar модельдерін жылдамдаттық. Бұл тәсіл бастапқы модельді нөлден бастап оқытуды талап етпесе де, шағын модель әлі де KV кэшінің айтарлықтай сыйымдылығын пайдаланады және TTFT көрсеткішін арттыра отырып, шағын префилл үстеме шығынын тудырады.

Бұл схема бойынша декодер префилл кезінде немесе аралас префилл-декодтау топтамаларында стандартты іріктеу арқылы токендерді генерациялай отырып, тек декодтауға арналған топтамаларға ғана болжам жасайды. Префилл кезеңінде мақсатты модельдің KV кэшінде жаңадан генерацияланған токенді де префилл жасау үшін мақсатты логиттер дереу іріктеледі. Бастапқы модель әлі іріктелмейді, бірақ ол шығаратын логиттер декодтау кезеңіне беріледі.

Ұлғайту

Декодтауда бастапқы модель алға жылжып, әр кезеңде ең жоғары токенді іріктейді. Қажетті бастапқы ұзындыққа жеткеннен кейін, токендер іріктеуші қабылданған тізбек ұзындығын анықтайтын логиттерді шығару үшін мақсатты модель арқылы өтеді. Қабылдау бастапқы модель мен мақсатты модельдің толық ықтималдық бөлінулерін салыстыру арқылы анықталады. Мақсатты модель қабылданған бастапқы тізбектен кейін әрқашан логиттердің бір жиынтығын шығаратындықтан, бұл қосымша шығысты өндіру үшін іріктеледі. Бастапқы модель бұл қабылданған токенді әлі көрмегендіктен, ол келесі декодтау қадамына дайындық ретінде өзінің KV кэш жазбаларын толтыру үшін қайта іске қосылады және логиттерді қайтадан тасымалдайды.

EAGLE

EAGLE — бұл ықтимал бастапқы токендерді ағаш тәрізді аралау арқылы генерацияланған бірнеше бастапқы тізбектерді зерттейтін спекулятивті кодтан шығару схемасы. Сызықтық тізбектегі ең жоғары балл жинаған токенді қадағалаудың орнына, әр түйіндік нүктеде Top-K кандидаттарын ескере отырып, бастапқы токендерді қатарынан орындау арқылы бекітілген (EAGLE) немесе динамикалық түрде пішімделген (EAGLE-2) ағаш зерттеледі. Содан кейін тізбектер бағаланады және мақсатты модельден қосымша токенді қоса отырып, жалғастыру үшін ең қолайлы ұзын тізбек таңдалады.

Ұлғайту

Дәлірек болжауға қол жеткізу үшін EAGLE бастапқы моделі тек токендерге ғана емес, сонымен қатар мақсатты модельдің мақсатты мүмкіндіктерін (соңғы қабаттың жасырын күйлерін) пайдалана отырып болжам жасайды. EAGLE кемшілігі — төмен кідіріс бюджеті шеңберінде қолайлы кандидаттарды генерациялауға жеткілікті дәл болатын реттелетін, шағын бастапқы модельдерді оқыту қажеттілігі. Әдетте, бастапқы модель бастапқы модельдің декодер қабатымен бірдей бір трансформерлік қабат болып табылады, ол өзінің қабаттарына және lm_head проекцияларына байланыстыру арқылы мақсатты модельмен тығыз біріктірілген. Бұл KV кэшінің аз сыйымдылығын талап ететіндіктен, EAGLE жадыны аз тұтынады.

Мақсатты модельдегі ағаш тәрізді тізбектерді тексеру үшін реттелетін назар аудару маскалары (attention masks) пайдаланылуы тиіс. Өкінішке орай, бүкіл тізбек үшін реттелетін назар аудару маскасын пайдалану шынайы кіріс ұзындықтары үшін назар аударуды едәуір баяулатады (50%-ға дейін), бұл болжам арқылы қол жеткізуге болатын жылдамдықтың бір бөлігін жоққа шығарады. Осы себепті біз әлі толық ағаш зерттеуін өндіріске енгізген жоқпыз, оның орнына DeepSeek-V3 техникалық есебінде ұсынылған MTP тәрізді схемалар арқылы бір токенді болжаудың арнайы жағдайына назар аудардық.

МТР (MTP)

Бұл схема болжау үшін токендермен қатар жасырын күйлердің қолданылуын қоспағанда, бастапқы-мақсатты декодтауға ұқсас. Тұрақты бастапқы-мақсатты болжаммен салыстырғанда префилл және декодтау кезеңдерінде де сәл көбірек жұмыс атқарылуы тиіс. Бастапқы модель токендерді де, жасырын күйлерді де пайдаланады: t_{i+1} токені t_i токеніне сәйкес келетін L_i логиттерінен іріктеледі, бұл өз кезегінде H_i жасырын күйлерінен алынады. Демек, кіріс токен буферлері мақсатты модель шығарған жасырын күй векторларына қатысты бір қадам солға жылжытылуы тиіс. Төмендегі суретте оқыту үшін қолданылатын сәйкестіктер, сондай-ақ инференс кезіндегі жылжу белгіленген.

Ұлғайту

Декодтау ағыны жасырын күйлер мен логиттердің екеуі де тасымалдануын қоспағанда, бастапқы-мақсатты декодтауға өте ұқсас. Біздің іске асыруымыз барлық байланысты іріктеу және логиттерді өңдеу логикасын бөліседі, тек модельдің тура шақыруларын ғана арнайы өңдейді. Бірнеше токен болжалған кезде, бастапқы модель болжау үшін бастапқы жасырын күйлерді пайдаланады, сондай-ақ өзінің мүмкіндіктері негізінде KV кэш жазбаларын толтырады. Ұзақ мерзімді перспективада бұл дәлдікті нашарлатуы мүмкін. Кейіннен мақсатты болжам үшін KV кэш жазбасын толтыру үшін бастапқы модельді іске қосқан кезде, біз оны дәлірек мақсатты жасырын күйлерді кіріс ретінде ала отырып, бүкіл тізбек бойынша іске қосамыз. Бұл бастапқы модельдер шағын болғандықтан, қосымша токендерді өңдеудің қосымша шығыны елеусіз болады.

MTP бастарын оқыту

MTP артықшылықтарын пайдалану үшін біз Perplexity деректер жиынтығында бапталған модельдерімізге қосылған MTP бастарын оқытуға қажетті инфрақұрылымды құрдық, ол 8xH100 құрылғылары бар бір түйіндік жүйеде жұмыс істейді. Шамамен бір күн ішінде біз Llama-1B-ден Llama-70B-ге дейін және DeepSeek V2-Lite модельдері үшін бастарды құра аламыз. Ірірек модельдер үшін біз баптау процесінде жасалған MTP бастарына сүйенеміз.

MTP оқытудың мақсаты — мақсатты жасырын күйлерден экстраполяцияланған бастапқы жасырын күйлер мен логиттерді мақсатты модельдің келесі токен логиттері мен жасырын күйлеріне сәйкестендіру. Жасырын күйлер үшін инференс қымбат тұратындықтан, біз оларды оқыту кезінде пайдалану үшін мақсатты модельдің инференс үшін оңтайландырылған іске асырылуы арқылы алдын ала есептейміз. Алайда, MTP инференс іске асырылуын тексеру және кванттау немесе оңтайландыру салдарынан сандық айырмашылықтардың нәтижелерге кедергі келтірмеуін қамтамасыз ету үшін тексеру шығыны мен дәлдікті бағалау үшін мақсатты модельдің де, бастапқы модельдің де инференс іске асырылуы толығымен қайта пайдаланылады.

Бастапқы мақалада қолданылған ShareGPT деректер жиынтығынан ірірек үлгілерге ауысқан кезде, EAGLE мақаласында көрсетілген және іске асырылған MTP бас архитектурасы 70B өлшемді модельдер үшін оқытылмай қалғанын байқадық. Қысқа тізбектердің үлкен санын қамтитын ShareGPT-тен айырмашылығы, біз тым ұзын емес, бірақ айтарлықтай үлкен промоунттардың азырақ санында оқытамыз. Бастапқы EAGLE бас өлшемдері әдеттегі трансформерден құрылымы жағынан сәл өзгеше болғандықтан, біз жойылған кейбір RMS қалыптастыру (Normalization) қабаттарын қайта енгіздік. Бұл оқытудың түйісуіне ғана емес, сонымен қатар бас дәлдігін бірнеше пайыздық тармаққа арттыруға мүмкіндік берді.

Ұлғайту

Қабат нормалары оқытуды жеңілдетіп қана қоймайды, нормаларды қайта енгізу математикалық тұрғыдан да түсінікті. MTP бас көлемі айтарлықтай үлкен болуы мүмкін болғандықтан (Llama 70B үшін шамамен 2 ГБ), мақсатты модельдің кірістірулері мен логит проекцияларын қайта пайдаланады. Оқыту кезінде бұлар қатырылады және MTP қабаты болжамдарды бастапқы модельдің проекциялық қабаты оқыту кезінде үйренген векторлық кеңістікке енгізуді үйренеді деп күтіледі. Нормаларды алып тастау арқылы жалғыз MLP нормадан кейінгі MLP сияқты функцияны үйренеді деп күтіледі, бұл бастапқы және мақсатты модельдердің жасырын күйлері арасындағы сәйкестікке кедергі келтіреді.

Спекулятивті кодтан шығару арқылы инференс

Инференс қозғалтқышында кіріс тізбектері үшін токендерді генерациялау мақсатында олар алдымен қолайлы мөлшердегі топтамаларға топтастырылуы тиіс, содан кейін келесі токендер үшін KV кэшінде беттер бөлінуі керек. Одан кейін кіріс токендері мен KV бетінің ақпараты модельді іске қосатын барлық параллель рангтарға таратылатын буферге жинақталады. Соңында, метадеректер GPU жадына көшіріледі және келесі токен іріктелетін логиттерді шығару үшін модель орындалады.

Олардың арасындағы сұрауларды үйлестіретін орама (wrapper) арқылы бастапқы және мақсатты инференс серверін еркін байланыстыратын кейбір іске асырулардан айырмашылығы, біздің бастапқы-мақсатты жұптарымыз тығыз байланысқан және генерациялау арқылы бірдей қадам жасайды. Топтаманы жоспарлау және KV бетін бөлу спекулятивті кодтан шығарудың барлық түрлері үшін модельдер арасында ортақ болып табылады: бұл модельді барлық қолданбалы интерфейсті көрсететін жан-жақты инференс серверімен байланыстыратын логиканы біріктіреді.

Perplexity жүйесіндегі инференс орындалу уақыты FlashInfer құралына негізделген, ол назар аудару ядросын конфигурациялау және жоспарлау үшін құрылуы тиіс метадеректерді анықтайды. Топтаманы құрайтын кейбір кіріс тізбектері үшін префилл, декодтау немесе тексеру үшін аралық буферлерді бөлу және назар аударуда қолданылатын белгілі бір тұрақты буферлерді толтыру үшін CPU жағындағы жұмыс орындалуы тиіс. Бұл жұмыс GPU пайдалануды барынша арттыру үшін жасырылуы тиіс кідірістерді тудыратын топтаманы жоспарлау және KV бетін бөлу шығындарынан басқа атқарылады.

Біз спецификациясыз инференс үшін CPU және GPU жағындағы жұмысты толығымен параллель еткенімізбен, спекулятивті кодтан шығару үшін CPU-GPU тепе-теңдігі күрделірек екенін анықтадық. Негізгі қиындық қабылданған токендер саны кейінгі іске қосу үшін тізбек ұзындығын анықтайтындығынан туындайды, бұл болдырмауға қиын GPU-CPU синхрондау нүктесін енгізеді. CPU жұмысының кідірісін жақсырақ жасыру үшін біз әртүрлі жоспарлау схемаларын тәжірибеден өткіздік.

Бастапқы-мақсатты кесте

Мақсатты модельден кішірек болғанымен, толық LLM бастапқы модель ретінде пайдаланылған кезде, ол әлі де GPU-да айтарлықтай кідіріс тудырады және қымбат CPU операцияларын жасыру үшін белгілі бір мүмкіндік береді. Шағын модельдер тензорлық параллелизмнен пайда көрмейтіндіктен, мақсат пен бастапқы модель бөлінетін рангтар саны арасында сәйкессіздік бар. Біздің іске асыруымызда бастапқы модель TP тобының негізгі рангінде ғана жұмыс істейді.

Ұлғайту

Бұрын айтылғандай, декодтау қадамы логиттерді келесі іске қосуға тасымалдайды. Бұл бастапқы модельдің бір рет орындалуын CPU жағындағы топтаманы жоспарлау жұмысымен қабаттастыруға мүмкіндік береді. Топтама біріктірілгеннен кейін, іріктеуші мен бастапқы модельге қайталанатын шақырулар бастапқы токендерді шығарады. Параллель түрде тексеруге арналған топтама мақсатты модель үшін біріктіріліп, параллель жұмысшылармен синхрондалады. Қабылданған тізбек ұзындығын анықтау үшін мақсатты логиттер тексеріліп, іріктеледі. Бұл сәтте келесі тізбек ұзындығын анықтау үшін GPU-CPU синхрондауы қажет. Бастапқы модель тек негізгі түйіндік нүктеде жұмыс істейтіндіктен, оның топтамасы ретті түрде орнатылады және мақсатты модель шығарған қосымша токенмен оның KV кэш жазбаларын толтыру үшін оның орындалуы басталады. Ағымдағы іске қосу кезінде осы бастапқы іске қосу арқылы шығарылған логиттер келесі іске қосу кезінде бірінші бастапқы токенді іріктеу үшін пайдаланылады. Ең мастығы, бастапқы модель жұмыс ілек жатқанда, келесі топтаманы жоспарлауға болады.

Бір токенге арналған MTP кестесі

Орындалу уақыты әлі де Eagle стиліндегі бастапқы ағаш зерттеуін қамтамасыз етпесе де, біз бір трансформерлік декодер қабатының өлшеміндегі модель шығарған бастапқы токендердің сызықтық тізбегін қарастыра отырып, осы схеманың арнайы жағдайын іске асырдық. Бұл схема DeepSeek R1 ашық бастапқы кодты салмақтарын пайдалана отырып, бастапқы болжау үшін қолданылуы мүмкін. Бір токенді болжаудың қосалқы жағдайы қызықты, өйткені ірі MTP қабаттары өздерінің үстеме шығындарын ақтау үшін жеткілікті жоғары қабылдау мөлшерлемелеріне қол жеткізеді.

Бастапқы модель әлдеқайда жылдам болып, CPU жағындағы кідірісті азырақ жасыратындықтан, MTP жоспарлауы сәл күрделірек. Сонымен қатар, бастапқы модель мақсатты модельмен бірге бөлінеді, бұл топтама туралы ақпарат үшін ортақ жады тасымалдауын талап етеді. Іске қосу топтама туралы ақпаратты тасымалдау және префилл логиттерінен бірінші токенді іріктеу арқылы басталады, бұл алдыңғы схемаға ұқсас. Келесі кезеңде D декодтау топтамасының өлшемі болып табылатын 2 * D токендерін өңдей отырып, токендерді тексеру үшін мақсатты модель іске қосылады. Бұл топтама екі жартыға тең бөлінетін болғандықтан, InfiniBand сияқты баяуырақ өзара байланыстар арқылы Mixture-of-Experts (MoE) модельдерінде шағын топтамаларға (micro-batching) өте қолайлы. Мақсатты модельдің жасырын күйлері келесі бастапқы іске қосуға тасымалданады, ал логиттер тексеру үшін іріктеушіге беріледі.

Ұлғайту

GPU-да қосымша жұмыстың шектеулі мөлшерін орындау арқылы біз бастапқы тізбек қабылданғаннан кейін GPU-CPU синхрондауынан аулақ боламыз. Мақсатты модельдің кіріс токендері жылжытылғаннан кейін, ядро келесі мақсатты токендерді тиісті орындарына қосады. Содан кейін бастапқы модель мақсатты модельмен бірдей топтама ақпаратымен қайта іске қосылады, KV кэш жазбаларын толтырады және келесі іске қосу үшін логиттер мен жасырын күйлерді құрады, қабылданбаған токендер бойынша кейбір артық жұмыстарды орындайды. Бұл жағдайларда пайдаланылмаған жұмыстың кідірісі бастапқы модельдің шағын өлшеміне байланысты шамамен байқалады. Бастапқы модельдің жұмысымен қатар тізбек ұзындығы CPU-да анықталады және GPU жұмысының аяқталуын күтпей-ақ келесі топтаманы жоспарлау басталады.

Бастапқы қабаттағы қосымша жұмыстың үстеме шығыны назар аудару кезінде байқалмайды, бірақ MLP қабаттары көбірек мәселе тудырады. Матрицалық көбейту нұсқаулары токендер саны өлшемі бойынша 64 шекарасына дейін толтырылатындықтан, егер екі еселеу айтарлықтай көбірек блоктарды талап етпесе, үстеме шығын жасырылады. Ұзағырақ бастапқы тізбектер үшін үстеме шығын қымбатырақ болады және тұрақты бастапқы-мақсатты модельдер үшін қолданылатын схема жақсырақ жұмыс істейді.

Әдебиеттер тізімі

Спекулятивті кодтан шығару арқылы трансформерлерден жылдам инференс

EAGLE: Спекулятивті іріктеу мүмкіндіктердің белсіздігін қайта қарауды талап етеді

EAGLE-2: Динамикалық бастапқы ағаштар арқылы тілдік модельдердің жылдамырақ инференсі

EAGLE-3: Оқыту кезіндегі тест арқылы үлкен тілдік модельдердің инференс жылдамдығын арттыру

Medusa: Бірнеше декодтау бастары бар қарапайым LLM инференсін жылдамдату құрылымы

Қабатты қалыптастыру (Layer Normalization)

FlashInfer: LLM инференсіне арналған тиімді және реттелетін назар аудару қозғалтқышы

DeepSeek-V3 техникалық есебі

Дәйексөз

@online{perplexity2025AcceleratingSonarThroughSpeculation,
  author = {{Perplexity} and Nandor Licker and Erin McKnight},
  title = {{Accelerating Sonar through speculation}},
  date = {2025-06-10},
  url = {https://perplexity.coolkeypoint.com/hub/blog/accelerating-sonar-through-speculation}
}