Az i14 Laplace 3B modellje fix memóriával fut AWS Trainium2-n
Az i14 Laplace 3B modellje a hagyományos transzformeres figyelem helyett lineáris figyelmet használ, ezért a következtetés során a rekurrens állapot mérete nem nő a kontextussal. Az AWS és az i14 a modellt natív PyTorch-támogatással futtatta Trainium2 gyorsítókon.
- A Laplace 3B GLA-alapú lineáris figyelmet használ.
- A rekurrens állapot szekvenciánként 29,5 MB, a kontextus hosszától függetlenül.
- A modell 3,028 milliárd paraméterből és 24 dekóderrétegből áll.
- Egy trn2.48xlarge példányon a tanítási sebesség körülbelül 120 000 token másodpercenként.
- 1280 párhuzamos szekvenciánál 11 544 kimeneti token volt a mért másodpercenkénti teljesítmény.
A KV-gyorsítótár helyett állandó méretű állapot
A legtöbb nagy nyelvi modell softmax-alapú figyelmet használ. Következtetéskor ezeknél a kulcs-érték gyorsítótár, vagyis a KV cache, a beszélgetés vagy a feldolgozott dokumentum kontextusával együtt növekszik. Ez csökkentheti a párhuzamosan kiszolgálható sorozatok számát, növelheti a kiszolgálás költségét, és az ügynökalapú alkalmazásokat arra kényszerítheti, hogy elvessenek korábbi előzményeket.
Az i14 Laplace 3B modellje ezzel szemben lineáris figyelmet alkalmaz. A korábbi kontextust a Gated Linear Attention, röviden GLA, egy rögzített méretű mátrixállapotba tömöríti. Ez az állapot szekvenciánként 29,5 MB, a kontextus hosszától függetlenül.
A vállalat modellje a figyelmi számítás tanítási skálázását O(T²) helyett O(T) értékre változtatja. Következtetéskor a modell minden tokennél frissíti a rekurrens állapotot, amelynek mérete nem változik a szekvencia hosszával. Az AWS szerint ez előnyös lehet hosszú ideig futó ügynökök, nagy dokumentumok elemzése és kiszámítható kiszolgálási igényű alkalmazások esetén.
A rögzített állapot ugyanakkor tömöríti az előzményeket, ezért az AWS külön kiemeli, hogy az információ-visszakeresés minőségét az adott felhasználási feladatokon kell értékelni.
A Laplace 3B felépítése
A Laplace 3B az i14 24 rétegű dekódere. Minden rétegben a softmax figyelmet Simple GLA váltja ki. A modell rejtett dimenziója 3072, 16 figyelmi fejet használ, a teljes paraméterszám pedig 3,028 milliárd. Ebből 2,718 milliárd paramétert Muon, 309,6 milliót pedig AdamW optimalizálóhoz rendeltek.
A GLA minden fejben egy 96 × 192 méretű állapotmátrixot tart fenn. Ez a korábbi kulcs-érték párok tömörített, fokozatosan halványuló rekordja. A fejek eltérő memóriatartományokkal dolgoznak: a decay ütemezés minimális értéke 1, maximális értéke 2048 token. Az első fej egy tokenes felezési idővel rendelkezik, a tizenhatodik fej felezési ideje megközelítőleg 2048 token.
A modell rétegenként egy négy token szélességű, oksági Conv1d konvolúciót, SiLU aktivációt és 8192 belső dimenziójú GLU-alapú előrecsatolt hálózatot használ. A tanításhoz 100 milliárd FineWeb-Edu tokent, 2048-as szekvenciahosszt és lépésenként 524 288 tokenes kötegméretet határoztak meg.
Három módosítás kellett a Trainium2-höz
A Laplace 3B eredeti megvalósítása GPU-kra optimalizált Triton kerneleket használt. Az AWS Generative AI Innovation Center csapata először standard PyTorch-műveletekkel írta újra a GLA-kernelt Trainium2-re. A két megközelítés összevetése után a béta állapotú TorchNeuron Native mellett döntöttek. Ez az AWS Trainium natív PyTorch-háttere.
Az első módosítás során a GPU-specifikus fused_chunk_simple_gla Triton-kernelt standard PyTorch-műveletekre cserélték. A számítást 16, egyenként 128 tokenes részre bontották, hogy illeszkedjen a NeuronCore lapkán található puffer 128 partíciójához.
A második változtatás a Muon kommunikációját érintette. A MuonFSDP2 két all_to_all kollektív művelete 24-szeresére csökkentette az egy rangra jutó adatforgalmat, és 1,63-szorosára növelte az áteresztőképességet. Harmadikként a szekvenciális visszaterjesztést statikusan kibontották, vagyis a GLA 16 részre bontott lépéseit rögzített gráffá alakították a fordító gráfméretre és lapkán belüli memóriára vonatkozó korlátai miatt.
Mérési eredmények egy trn2.48xlarge példányon
A bemutatott konfigurációban egy trn2.48xlarge példány 16 Trainium2 lapkát, 64 logikai NeuronCore magot, nagy sávszélességű memóriát és NeuronLink-v3 összeköttetést egyesít. A Trainium2 megvalósítását a szoftveresen kezelt lapkán belüli memória, a lapkák közötti gyors kapcsolat és a statikus fordító egyaránt meghatározta.
Az i14 és az AWS közlése szerint a Laplace 3B ezen a példányon körülbelül 120 000 tokent tanított másodpercenként. Kiszolgáláskor 1280 párhuzamos szekvencia mellett 11 544 kimeneti token volt a mért másodpercenkénti teljesítmény.
A bejelentés jelentősége a felhasználók számára az lehet, hogy a Laplace 3B-nél a kiszolgálási állapot mérete nem nő minden új tokennel. Ez azonban a lineáris figyelem sajátos kompromisszumával jár: a korábbi kontextus tömörített állapotban marad, ezért az adott alkalmazás visszakeresési feladatait külön kell tesztelni.
