Urychlení modelu Sonar prostřednictvím spekulace

Spekulativní dekódování zvyšuje rychlost generování velkých jazykových modelů (LLM) pomocí rychlého a malého návrhového modelu, který vytváří kandidáty na dokončení ověřované větším cílovým modelem. V tomto schématu se namísto jednoho spuštění drahého cíle, které vytvoří jediný token, v jednom kroku vygeneruje více tokenů. Zde představujeme podrobnosti implementace různých druhů spekulativního dekódování použitých v Perplexity ke snížení latence mezi tokeny u modelů Sonar.

Spekulativní dekódování

Spekulativní dekódování využívá strukturu přirozených jazyků a autoregresní povahu transformerů k urychlení generování tokenů. Ačkoli větší modely, jako je Llama-70B, nesou více znalostí než menší, jako je Llama-1B, v některých jednodušších úlohách fungují podobně. Tento překryv naznačuje, že určité sekvence jsou lépe generovány méně nákladnými modely, přičemž složité problémy zůstávají na bedrech větších modelů. Výzva spočívá v určení, která dokončení jsou lepší a zda je generování menšího modelu stejně kvalitní jako generování většího.

Naštěstí jsou modely LLM autoregresní transformery: když dostanou sekvenci tokenů, vydají pravděpodobnostní rozdělení dalšího tokenu. Logity odvozené z mezilehlých funkcí spojených s tokeny ve vstupní sekvenci navíc ukazují, jak je pro model pravděpodobné, že vydá právě tyto tokeny. Tato vlastnost umožňuje spekulaci: pokud je sekvence tokenů generována menším modelem počínaje vstupní předponou, lze ji spustit přes větší model, aby se zjistilo, jak dobře se shoduje s cílovým modelem. Každá předpona kandidátů je ohodnocena pravděpodobností a vybere se nejdelší nad prahem přijetí. Jako bonus poskytuje cílový model také následný token zdarma: pokud model návrhu vygeneruje n tokenů, v jednom kroku lze vydat až n + 1.

Zvětšit

V době odvozování lze proces spekulativního vzorkování rozdělit přibližně do 4 fází:

Předběžné vyplnění (prefill): cílový i návrhový model musí být spuštěny na vstupní sekvenci, aby se naplnily položky mezipaměti KV. Zatímco některá schémata, jako je Medusa, používají pro predikci jednodušší husté vrstvy, v tomto příspěvku se zaměřujeme na návrhy založené na transformerech, které potřebují vlastní mezipaměti KV.

Generování návrhu: model návrhu iteruje a vytváří řadu fixních tokenů. Sekvence návrhu může být lineární nebo může model prozkoumat stromovou strukturu až do dané hloubky (EAGLE, Medusa). Zde se zaměřujeme na lineární sekvence.

Přijetí: cílový model běží na sekvenci návrhu a buduje logity odpovídající každému tokenu návrhu. Určuje se délka nejdelší přijatelné sekvence.

Cílové generování: jelikož cílový model generoval logity, v nesouladné pozici nebo na konci sekvence odpovídají logity následnému tokenu. Tyto logity lze vzorkovat, aby poskytly robustní token z cíle, čímž se sekvence uzavře.

Existují různé metody implementace spekulativního dekódování. V tomto příspěvku se zaměříme na schémata, která jsme použili k urychlení modelů Sonar pomocí interního 1B modelu, a také na predikční mechanismy, které budujeme pro urychlení modelů v měřítku DeepSeek.

Cíl-Návrh

Spekulativního dekódování lze dosáhnout propojením stávajícího malého LLM jako modelu návrhu s cílovým modelem za účelem generování kandidátních sekvencí. V produkčním prostředí jsme urychlili Sonar pomocí modelu Llama-1B doladěného na stejné datové sadě jako cíl. Ačkoli tento přístup nevyžadoval trénování návrhu od nuly, malý model přesto využívá významnou kapacitu mezipaměti KV a zavádí mírnou režii předběžného vyplnění, čímž zvyšuje TTFT.

V rámci tohoto schématu dekodér spekuluje pouze u dávek určených pouze k dekódování a generuje tokeny prostřednictvím standardního vzorkování během předběžného vyplnění nebo u smíšených dávek předběžného vyplnění a dekódování. Ve fázi předběžného vyplnění jsou cílové logity okamžitě vzorkovány, aby se předběžně vyplnil také nově vygenerovaný token v mezipaměti KV modelu návrhu. Model návrhu ještě není vzorkován, ale logity, které produkuje, se přenášejí do fáze dekódování.

Zvětšit

Při dekódování je model návrhu posunut vpřed a v každé fázi vzorkuje špičkový token. Po dosažení požadované délky návrhu jsou tokeny prohnány cílovým modelem, aby se vytvořily logity, na jejichž základě vzorkovač identifikuje délku přijaté sekvence. Přijetí se určuje porovnáním úplných pravděpodobnostních rozdělení z návrhu a cíle. Jelikož cíl vždy vydá jednu sadu logitů po přijaté sekvenci návrhu, je tato sada vzorkována, aby se vytvořil dodatečný výstup. Jelikož model návrhu tento přijatý token ještě neviděl, je spuštěn znovu, aby naplnil své odpovídající položky mezipaměti KV jako přípravu na další krok dekódování a opět přenesl logity.

EAGLE

EAGLE je schéma spekulativního dekódování, které zkoumá více sekvencí návrhu generovaných prostřednictvím stromového procházení pravděpodobných tokenů návrhu. Pevný (EAGLE) nebo dynamicky tvarovaný (EAGLE-2) strom je zkoumán pomocí po sobě jdoucích spuštění tokenů návrhu, přičemž se v každém uzlu zvažují kandidáti Top-K namísto následování tokenu s nejvyšším skóre v lineární sekvenci. Sekvence jsou poté ohodnoceny a vybere se nejdelší vhodná pro pokračování, přičemž se připojuje další token z cíle.

Zvětšit

Aby bylo možné dosáhnout přesnější predikce, model návrhu EAGLE predikuje nejen na základě tokenů, ale využívá také cílové rysy (skryté stavy poslední vrstvy) cílového modelu. Nevýhodou EAGLE je potřeba trénovat vlastní, malé modely návrhu, které jsou dostatečně přesné pro generování vhodných kandidátů v rámci nízkého rozpočtu na latenci. Model návrhu je obvykle jediná vrstva transformeru totožná s vrstvou dekodéru původního modelu, která je pevně spjata s cílem pomocí propojení s jeho vnořeními a projekcemi lm_head. Jelikož to vyžaduje menší kapacitu mezipaměti KV, má EAGLE nižší paměťovou náročnost.

K ověření stromových sekvencí v cílovém modelu je nutné použít vlastní masky pozornosti. Bohužel použití vlastní masky pozornosti pro celou sekvenci výrazně zpomaluje pozornost u realistických délek vstupů (až o 50 %), což neguje část zrychlení dosažitelného spekulací. Z tohoto důvodu jsme dosud plný průzkum stromu do produkčního prostředí nasadili, přičemž se namísto toho soustředíme na speciální případ predikce jednoho tokenu prostřednictvím schémat podobných MTP, která jsou prezentována v technické zprávě DeepSeek-V3.

MTP

Toto schéma je podobné dekódování návrhu a cíle s tou výjimkou, že se pro predikci používají skryté stavy spolu s tokeny. Ve fázích předběžného vyplnění i dekódování je třeba odvést o něco více práce ve srovnání s běžnou spekulací návrhu a cíle. Model návrhu používá tokeny i skryté stavy: token t_{i+1} je vzorkován z logitů L_i odpovídajících tokenu t_i, které jsou zase odvozeny ze skrytých stavů H_i. V důsledku toho musí být vstupní vyrovnávací paměti tokenů posunuty o jeden krok doleva vzhledem k vektorům skrytých stavů vydaným cílem. Následující obrázek označuje shody použité pro trénování a také posun během odvozování.

Zvětšit

Tok dekódování je velmi podobný dekódování návrhu a cíle s tou výjimkou, že se přenášejí jak skryté stavy, tak logity. Naše implementace sdílí veškerá související vzorkování a zpracování logitů a specializuje pouze vyvolání dopředného průchodu modelu. Když je předpovězeno více tokenů, model návrhu používá k predikci skryté stavy návrhu a také naplňuje položky mezipaměti KV na základě vlastních rysů. Z dlouhodobého hlediska to může snížit přesnost. Následně při spuštění modelu návrhu za účelem naplnění položky mezipaměti KV pro cílovou predikci jej spustíme na celé sekvenci a jako vstupy bereme přesnější skryté stavy cíle. Jelikož jsou tyto modely návrhu malé, dodatečné náklady na zpracování dalších tokenů jsou zanedbatelné.

Trénování hlav MTP

Abychom mohli těžit z MTP, vytvořili jsme infrastrukturu potřebnou k trénování hlav MTP připojených k našim doladěným modelům na datových sadách Perplexity, které běží na jednom uzlu s 8 zařízeními H100. Za přibližně jeden den dokážeme vytvořit hlavy pro modely od Llama-1B po Llama-70B a DeepSeek V2-Lite. U větších modelů spoléháme na hlavy MTP vytvořené během procesu doladění.

Cílem trénování MTP je sladit skryté stavy návrhu a logity extrapolované z cizích skrytých stavů s logity dalšího tokenu a skrytými stavy cíle. Jelikož je odvozování pro skryté stavy nákladné, předem je vypočítáme pomocí naší implementace cílového modelu optimalizované pro odvozování, která se použije během trénování. K ověření ztráty při ověřování a odhadu přesnosti však plně znovu používáme implementaci odvozování jak cílového, tak návrhového modelu, abychom ověřili implementaci MTP pro odvozování a zajistili, že číselné rozdíly způsobené kvantováním nebo optimalizacemi nebudou bránit výsledkům.

Při škálování z datové sady ShareGPT použité v původním článku na větší vzorky jsme si všimli, že architektura hlavy MTP popsaná a implementovaná v článku EAGLE se u modelů velikosti 70B nepodařilo trénovat. Na rozdíl od ShareGPT, která obsahovala větší počet kratších sekvencí, trénujeme na o něco menším počtu podstatně delších výzev. Jelikož se původní hlavy EAGLE strukturou mírně odchylovaly od typického transformeru, znovu jsme zaviedli některé vrstvy normalizace RMS, které byly odstraněny. Zjistili jsme, že to nejen umožnilo konvergenci trénování, ale také to zvýšilo přesnost hlav o několik procentních bodů.

Zvětšit

Normy vrstev nejen usnadňují trénování, ale jejich opětovné zavedení je také matematicky intuitivní. Hlavy MTP znovu používají vnoření a projekce logitů cílového modelu, protože mohou být značně velké (přibližně 2 GB pro Llama 70B). Během trénování jsou tyto zmrazeny a očekává se, že vrstva MTP se naučí vkládat predikce do stejného vektorového prostoru, jaký se projekční vrstva původního modelu naučila během trénování. Vynecháním norem se očekává, že jediná MLP se naučí stejnou funkci jako MLP následovaná normou, což brání sladění mezi skrytými stavy návrhu a cílovými modely.

Odvozování se spekulativním dekódováním

V odvozovacím stroji je nutné pro generování tokenů pro vstupní sekvence tyto nejprve seskupit do dávek přiměřené velikosti a poté přidělit stránky v mezipaměti KV pro další tokeny. Vstupní tokeny a informace o stránkách KV jsou poté zabaleny do vyrovnávací paměti vysílané do všech paralelních hodností spouštěcích model. Nakonec jsou metadata zkopírována do paměti GPU a model je spuštěn za účelem vytvoření logitů, ze kterých se vzorkuje další token.

Na rozdíl od některých implementací, které volně propojují server odvozování návrhu a cíle prostřednictvím obálky orchestrující požadavky mezi nimi, jsou naše páry návrhu a cíle pevně propojeny a procházejí generováním v jednotném rytmu. Plánování dávek a přidělování stránek KV je sdíleno mezi modely pro všechny formy spekulativního dekódování: to sjednocuje logiku propojující model s nadřazeným serverem odvozování, protože všechny vystavují stejné rozhraní.

Doba běhu odvozování v Perplexity je vytvořena kolem FlashInfer, což určuje metadata, která je třeba vytvořit pro konfiguraci a plánování jádra pozornosti. Vzhledem k některým vstupním sekvencím tvořícím dávku je pro předběžné vyplnění, dekódování nebo ověření nutné provést práci na straně CPU k přidělení mezipamětí a naplnění určitých konstantních vyrovnávacích pamětí používaných v pozornosti. Tato práce je dodatečná k nákladům na plánování dávek a přidělování stránek KV, které rovněž způsobují latence, jež je nutné skrýt pro maximalizaci využití GPU.

Zatímco jsme plně zparalelnili práci na straně CPU a GPU pro odvozování bez spekulace, zjistili jsme, že rovnováha mezi CPU a GPU pro spekulativní dekódování je složitější. Hlavní výzva vyplývá z toho, že počet přijatých tokenů určuje délku sekvence pro následné spuštění, což přináší bod synchronizace mezi GPU a CPU, kterému se lze jen těžko vyhnout. Experimentovali jsme s různými schématy plánování, abychom co nejlépe skryli latenci práce CPU.

Plán návrhu a cíle

Přestože je menší než cílový model, když je jako návrh použit celý LLM, stále způsobuje na GPU značnou latenci a poskytuje určitou rezervu pro skrytí drahých operací CPU. Vzhledem k tomu, že menší modely netěží z paralelismu tenzorů, existuje nesoulad mezi počtem hodností, mezi nimiž je cíl a návrh rozdělen (sharded). V naší implementaci běží model návrhu pouze na vedoucí hodnosti skupiny TP.

Zvětšit

Jak již bylo uvedeno, krok dekódování přenáší logity do dalšího spuštění. To nám umožňuje překrýt jedno spuštění modelu návrhu s prací plánování dávky na straně CPU. Po sestavení dávky vytvářejí opakovaná volání vzorkovače a návrhu tokeny návrhu. Paralelně se sestavuje dávka pro ověření pro cílový model a synchronizuje se s paralelními pracovníky. Cílové logity jsou ověřeny a vzorkovány pro určení délek přijatých sekvencí. V tomto bodě je nutná synchronizace GPU-CPU k určení následných délek sekvencí. Jelikož model návrhu běží pouze na vedoucím uzlu, jeho dávka je nastavena sekvenčně a jeho spuštění je zahájeno k naplnění položek mezipaměti KV dalším tokenem, který cíl vytvořil. Logity vytvořené tímto spuštěním návrhu v aktuálním spuštění budou použity k vzorkování prvního tokenu návrhu v následném spuštění. A co je nejdůležitější, během běhu návrhu lze naplánovat další dávku.

Plán MTP pro jeden token

Zatímco doba běhu dosud neposkytuje prozkoumání stromu návrhu ve stylu Eagle, implementovali jsme speciální případ tohoto schématu a uvažujeme lineární sekvenci tokenů návrhu vytvořenou modelem o velikosti jediné vrstvy dekodéru transformeru. Toto schéma lze použít pro predikci návrhu pomocí otevřených vah DeepSeek R1. Dílčí případ predikce jednoho tokenu je zajímavý, protože velké vrstvy MTP dosahují dostatečně vysokých měr přijetí, aby ospravedlnily svou režii.

Plánování MTP je poněkud složitější, protože model návrhu je mnohem rychlejší a skrývá menší latenci na straně CPU. Návrh je navíc rozdělen spolu s cílovým modelem, což vyžaduje přenosy sdílené paměti pro informace o dávce. Spuštění začíná přenosem informací o dávce a vzorkováním prvního tokenu z přenesených logitů, podobně jako v předchozím schématu. Dále je spuštěn cíl pro ověření tokenů, který zpracovává 2 * D tokenů, kde D je velikost dávky dekódování. To je ideální pro dávkování v modelech typu Mixture-of-Experts (MoE) přes pomalejší propojení, jako je InfiniBand, protože dávka se rovnoměrně rozdělí na dvě poloviny. Skryté stavy cíle se přenášejí do dalšího spuštění návrhu, zatímco logity jsou předány do vzorkovače k ověření.

Zvětšit

Provedením omezeného množství dodatečné práce na GPU se vyhýbáme synchronizaci mezi CPU a GPU po přijetí sekvence návrhu. Po posunutí vstupních tokenů cílů jádro zapojí další cílové tokeny do jejich odpovídajících míst. Návrh je pak spuštěn znovu se stejnými informacemi o dávce jako cíl, čímž se naplní položky mezipaměti KV a vytvoří se logity a skryté stavy pro další spuštění, přičemž se provádí určitá nadbytečná práce na tokenech, které nebyly přijaty. V těchto situacích je latence nevyužité práce sotva měřitelná kvůli malé velikosti modelu návrhu. Paralelně s během návrhu jsou na CPU určeny délky sekvencí a je zahájeno plánování další dávky, aniž by bylo nutné čekat na dokončení práce GPU.

Režie dodatečné práce ve vrstvě návrhu není v pozornosti znatelná, avšak vrstvy MLP jsou problematičtější. Jelikož instrukce násobení matic zarovnávají na hranici 64 podél dimenze počtu tokenů, pokud zdvojnásobení nevyžaduje podstatně více bloků, je režie skryta. Pro delší sekvence návrhu je režie dražší a schéma používané pro běžné modely návrhu a cíle funguje lépe.

Odkazy

Rychlé odvozování z transformerů prostřednictvím spekulativního dekódování

EAGLE: Spekulativní vzorkování vyžaduje přehodnocení nejistoty rysů

EAGLE-2: Rychlejší odvozování jazykových modelů s dynamickými stromy návrhu

EAGLE-3: Škálování urychlení odvozování velkých jazykových modelů prostřednictvím testu v čase trénování

Medusa: Jednoduchý rámec urychlení odvozování LLM s vícenásobnými dekódovacími hlavami

Normalizace vrstev

FlashInfer: Efektivní a přizpůsobitelný stroj pozornosti pro poskytování odvozování LLM

Technická zpráva DeepSeek-V3

Citace

@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}
}