Miksi jotkut tekoälymallit ovat nopeampia kuin toiset

Linkkitaulukko

Tiivistelmä ja 1 Johdanto

2. Aiheeseen liittyvät työt

3. Tausta

3.1. Merkintä ja 3.2. Kielimallin päättely

3.3. Multi-Query, Multi-Head ja yleistetty usean kyselyn huomio

4. Kontekstitietoinen kaksijakoinen huomio ja 4.1. Motivaatio

4.2. Formulaatio ja 4.3. Muistin IO monimutkaisuus

5. Kokeilut

5.1. Multi-Head-, Multi-Query- ja Multi-Group Attention -ominaisuuksien vertailu

5.2. Ominaisuuksia vastaavien mallien viiveet

5.3. Sovellukset

6. Johtopäätös ja viitteet

A. Usein kysytyt kysymykset

B. Aiheeseen liittyvä työ

C. Asennus

D. Multi-Group Attention Family

E. Kontekstitietoinen kaksijakoinen huomio

F. Sovellukset: lisätuloksia

G. Yhteensopivuus spekulatiivisen dekoodauksen ja nopean dekoodauksen tekniikoiden kanssa

2. Aiheeseen liittyvät työt

Kirjallisuudessa on useita tapoja parantaa päätelmien latenssia ja/tai latenssia. Kvantisointi vähentää muistin käyttöä käyttämällä matalan bittileveyden esityksiä, kuten int8, int4 ja fp8 (Wei et al., 2023; Yao ym., 2022; Dettmers et al., 2022; Frantar et al., 2022; Kuzmin et al., 202 et al., Xia2o02; Xia2 et al., 2022;). Kvantisointi, kun sitä käytetään vain malliparametreihin, tarjoaa pieneneviä tuloksia, kuten pitemmillä sekvenssipituuksilla ja suurilla eräkooilla, joissa muistin käyttö ja pistetuotteen huomioimiseen liittyvä laskenta hallitsevat yleistä päättelyviivettä.

Harva huomiota (Beltagy ym., 2020; Child et al., 2019; Zaheer ym., 2020) on tutkittu laajasti keinona vähentää huomion monimutkaisuutta pidemmissä konteksteissa ja nopeammissa päätelmissä. Pope et ai. (2022) tutkii suurten kielimallien generatiivisen päättelyn tehokkuutta käyttämällä moniulotteisia osiointitekniikoita, jotka on optimoitu TPU:ille (kollektiivinen einsum), jotta saavutetaan Pareto-raja latenssin ja mallin FLOP-käytön suhteen. Paperi osoittaa myös, että usean kyselyn huomioiminen mahdollistaa jopa 32x suuremman kontekstin pituuden skaalatauksen painottaen tehokkuutta suuressa eräkoossa. Sivutettu huomio (Kwon et al., 2023) parantaa KV-välimuistin muistinhallintaa jakamalla sen lohkoihin ja käyttämällä lohkotaulukkoa kartoitustarkoituksiin. Tämä lähestymistapa mukautuu tehokkaasti dynaamisiin työkuormituksen siirtoihin ja vähentää muistin tallennusvaatimuksia jakamalla kehotteen KV-välimuistin useiden tulostussarjojen kesken. Tämä ei kuitenkaan vähennä KV-välimuistin muistilukuja.

Spekulatiivinen dekoodaus ja sen muunnelmat käyttävät pienempää luonnosmallia ehdottaakseen useita peräkkäisiä tokeneita, joita päämalli käsittelee rinnakkain tällaisten tokenien hyväksymiseksi tai hylkäämiseksi (Chen et al., 2023; Leviathan et al., 2022; Li et al., 2024; Cai et al., Fu et al., 2024,2). Keskeisenä ideana on mahdollistaa useiden tokenien dekoodaus joka vaiheessa, jolloin päämallin muistin IO-käytöt kuoletetaan. Dekoodauksen latenssia hallitsee kuitenkin edelleen KV-välimuistin I/O-kaistanleveys suurilla kontekstikokoilla, jolloin kaksihaarainen huomio voi parantaa dekoodausnopeutta entisestään. Lyhyesti sanottuna, inkrementaalinen dekoodaus keskittyy pienentämään mallin latauksen kuoletetun muistin IO:ta, kun taas monikysely ja kaksihaarainen huomio alentavat KV-välimuistin muistin IO:ta.

3. Tausta

3.1. Merkintä

Käytämme koko paperissa seuraavaa merkintää.

3.2. Kielimallin päättely

Kielimallille on olemassa monia johtopäätösskenaarioita, mukaan lukien eräpäättely ja yhden kontekstin eränäyte (kuva 1). Eräpäätelmä viittaa tapaukseen, jossa käsittelemme useita syötteitä yhdessä erässä ja luomme seuraavat tunnukset jokaiselle eräindeksille itsenäisesti. Jos eräkoko on 1, tämä pelkistyy yhden kontekstin päättelyyn. Toinen skenaario on yhden kontekstin eränäyte, jossa luomme useita sekvenssejä yhden kontekstin perusteella, jossa ero eräpäättelytapauksen välillä on, että esitäyttö on tehtävä vain yhdelle kontekstille KV-välimuistin saamiseksi, minkä jälkeen se lähetetään muille eräindekseille.

Kuva 1 havainnollistaa myös kielimallin päättelyn kahta vaihetta: (a) kontekstin koodaus tai esitäyttö ja (b) inkrementaalinen dekoodaus. Kontekstikoodaus viittaa yhteen eteenpäin tapahtuvaan siirtoon, joka laskee avain- ja arvotensorit kaikille kontekstin merkkipaikoille. Kun avain- ja arvotensorit on laskettu, tallennamme nämä avain- ja arvotensorit välimuistiin käytettäväksi huomiomekanismissa inkrementaalisen dekoodausvaiheen aikana, joka luo peräkkäin yhden merkin kerrallaan[2].

Kontekstikoodausvaiheen aikana liukulukuoperaatioiden määrä suhteessa muistin sisääntulo/lähtö (IO) -toimintoihin on suuri, mikä vastaa laskentaan sidottua järjestelmää, jossa FLOP:t vaikuttavat latenssiin. Kuitenkin inkrementaalisen dekoodauksen aikana, jossa kiinnitämme huomiota yhdelle kyselytunnisteelle, tämä kuuluu muistiin sidottuun järjestelmään, jossa laskutoimitusten määrä muistiin pääsyä kohti on suunnilleen 1-1 (katso lisätietoja liitteestä D.1). Muistin IO viittaa luku- ja kirjoitustoimintoihin suuren kaistanleveyden muistista (HBM) (Jia et al., 2018) nopeaan on-chip SRAM:iin, jossa varsinainen laskenta tapahtuu. Itse inkrementaalisen dekoodauksen muisti IO koostuu kahdesta osasta: (1) malliparametrin lataus ja (2) KV-välimuistin lataus. Komponentti (1) on vakio riippumatta kontekstin pituudesta m tai erän koosta b, jossa komponentti (2) riippuu sekä m:stä että b:stä ja hallitsee kokonaismuistia IO, jos m tai b ovat korkeat, mikä voi muodostua merkittäväksi pullonkaulaksi johtopäätöksille. Työmme keskittyy ensisijaisesti komponentin (2) vähentämiseen.