Förtjänar ditt viktformat en egen kärna?

Du har en idé om hur modellvikter kan ta mindre plats än i MLX:s inbyggda format. MLX stöder redan tvåbitarskvantisering, men två bitar per vikt är inte hela lagringskostnaden. Kanske räcker tre möjliga värden per vikt, packade i en annan layout. Du kan rita upp dina byte på papper, men MLX har ingen operation som läser den layouten.

Innan du skriver den operationen finns en fråga värd att besvara: kan modellen fortfarande göra något användbart med de tal ditt recept bevarar?

Det finns ett experiment att göra först: packa vikterna, packa upp dem till vanliga flyttal igen och kör modellen. Körmiljön ser en välbekant checkpoint. Modellen får leva med de tal som ditt recept lämnade kvar.

Vi använder det här arbetssättet i vårt eget arbete med låga bitbredder för att jämföra recept innan vi lägger tid på en beräkningskärna, koden som utför beräkningen på processorn. Här provar vi det med NumPy och vanlig MLX och använder perplexitet för att mäta hur förmågan att förutsäga nästa token påverkas.

Som jämförelse är det här några av de viktrepresentationer som finns i vanlig MLX.¹

De affina raderna använder grupper om 64 med en 16-bitars skalfaktor och en 16-bitars bias. Det lägger till 0,5 bitar per vikt utöver koderna. Filhuvuden och vikter som lämnas oförändrade är inte medräknade.

Vår layout med tre värden, alltså en ternär layout, använder 1,875 bitar per utvald vikt vid samma gruppstorlek, inklusive skalfaktorn. MLX kan representera de värdena i sitt befintliga tvåbitarsformat, men kan inte köra vår mer kompakta layout direkt.

Vad överlever turen fram och tillbaka?

Kvantiseringen väljer talen. Flyttalsvikter blir en gemensam skalfaktor och små heltalskoder. Då går i regel information förlorad: flera ursprungliga värden får nu samma rekonstruerade värde.

Packningen ordnar koderna i byte. De valda koderna ska kunna lagras och återskapas utan förlust. Hur de är ordnade ska inte påverka modellkvaliteten.

Körmiljön räknar med representationen. Den kan till exempel avkoda små delar inne i en matrismultiplikation och därmed undvika en fullständig flyttalskopia.

Vårt experiment avkodar paketet till en checkpoint i bfloat16, eller bf16, ett 16-bitars flyttalsformat som modellen redan stöder. Att lägga den begränsade uppsättningen värden i ett rymligare format återställer inte informationen som kvantiseringen tog bort.

PyTorchs fake quantization bygger på den välbekanta idén att simulera kvantisering och avkvantisering med flyttal som utdata.² Genom att skriva och läsa det faktiska paketet testar vi också den lagringslayout vi föreslår.

Tre värden och en gemensam skalfaktor

I vårt exempel blir varje vikt ett av värdena $\{-s, 0, +s\}$. Tecknet eller nollan är en trit, bitens släkting i bas tre. En grupp med $G$ vikter i följd delar skalfaktorn $s$.

Vi använder en medvetet enkel regel: ta medelvärdet av vikternas absolutbelopp i gruppen som skalfaktor, avrunda den till fp16 för lagring och avrunda sedan de normaliserade vikterna och begränsa dem till intervallet:

$$s = \operatorname{fp16}\!\left(\frac{1}{G}\sum_{j=1}^{G}|W_j|\right)$$ $$t_j = \operatorname{clip}\!\left(\operatorname{round}(W_j/s), -1, 1\right)$$

Återskapa varje vikt som $\widehat{W}_j = \operatorname{bf16}(s\,t_j)$. Skalfaktorn använder fp16, ett annat 16-bitars flyttalsformat än bf16. Genom att räkna med den lagrade skalfaktorn får vi med även dess avrundning i experimentet. Avrundningen väljer det jämna alternativet när två tal ligger lika nära. Om skalfaktorn är noll återskapas en grupp med enbart nollor, utan division.

Regeln är lätt att granska. Om en modell som tränats med vanliga vikter tål att göras ternär i efterhand återstår att se.

Fem tritar i en byte

En byte har 256 möjliga värden. Fem tritar har $3^5 = 243$ möjliga kombinationer, så alla fem får plats. Avbilda tritarna $\{-1,0,+1\}$ på siffrorna $\{0,1,2\}$ och skriv ett vanligt tal i bas tre:

$$b = d_0 + 3d_1 + 9d_2 + 27d_3 + 81d_4$$

Det här lilla NumPy-exemplet packar och packar upp en byte. Den första siffran är den minst signifikanta:

import numpy as np

trits = np.array([-1, 0, 1, -1, 1], dtype=np.int16)
powers = np.array([1, 3, 9, 27, 81])

packed = np.uint8((trits + 1) @ powers)  # 183
written = bytes([int(packed)])
byte = int.from_bytes(written, "little")
recovered = (byte // powers) % 3 - 1

np.testing.assert_array_equal(recovered, trits)

Det fullständiga exemplet skriver grupper av koder och skalfaktorer till disk. Avkodaren får bara det paketet, ingen ursprunglig checkpoint som den skulle kunna återskapa de avsedda vikterna från.

Innan du frågar modellen något, testa alla 243 möjliga kombinationer av fem tritar mot en oberoende skalär beräkning. Testa också ofullständiga avslutande byte, gruppgränser och de oanvända bytevärdena 243 till 255. Vårt format avvisar de värdena och kräver att utfyllnadspositioner avkodas till noll.

En lyckad tur fram och tillbaka betyder att de valda tritarna överlevde. Om de var bra val är ett annat test.

Räkna de byte du faktiskt skrev

Fem tritar per byte antyder $8/5 = 1{,}6$ bitar per vikt. Vår layout börjar däremot varje grupp på en bytegräns, så en grupp med 64 vikter behöver 13 byte för koderna. Lägg till två byte för skalfaktorn i fp16, så upptar gruppen 15 byte.

För själva viktdata i den här layouten blir det alltså:

$$\text{bitar per vikt} = \frac{8\bigl(\lceil G/5\rceil + 2\bigr)}{G}$$

Siffrorna gäller nyttolasten för de utvalda vikterna, inklusive utfyllnad inom grupperna. Filhuvuden, tensorformer och vikter som lämnas oförändrade tar också plats. Exemplet redovisar hela paketets storlek i byte separat.

Ternära format med fler begränsningar kan använda ännu färre bitar för koderna. Sherry kräver till exempel exakt en nolla och tre tritar som inte är noll i varje block om fyra vikter. Det lämnar 32 möjliga mönster, som ryms i fem bitar: 1,25 kodbitar per vikt, innan skalfaktorer, metadata och utfyllnad räknas in.³ Besparingen kommer av att begränsa vilka viktmönster som tillåts, inte av att packa godtyckliga tritar tätare. Ett recept inspirerat av den begränsningen är ytterligare en kandidat att packa, packa upp och kvalitetstesta innan man lägger till stöd i körmiljön.

Vi kan mäta paketets filstorlek innan en kärna finns. Vad körningen i bf16 inte kan tala om är hur mycket minne en implementation som räknar direkt med formatet skulle behöva, eller hur snabbt den skulle kunna använda dessa byte.

Fråga modellen

Exemplet använder textdelen av Qwen3.5-0.8B:s MLX-checkpoint i bf16. Vi ändrar vikterna i feed-forward-lagrens gate-, up- och down-projektioner. Attention, embeddings, normalisering och övriga parametrar behåller sina ursprungliga värden. Experimentet gäller en del av modellen, inte ett påstående om att varje vikt kan lagras med bitåtgången ovan.

För varje gruppstorlek skriver skriptet ett paket, avkodar det till en ny bf16-checkpoint, laddar den på nytt genom vanlig MLX LM och utvärderar den mot den orörda referensen. Båda får exakt samma sparade tokenblock.

Perplexitet mäter hur mycket sannolikhet modellen tilldelar nästa token i en text: det är exponentialfunktionen av medelvärdet av de negativa log-sannolikheterna för dessa token. Lägre är bättre.

Vi använder 64 icke-överlappande block med 512 token ur valideringsdelen av engelska WikiText-2 raw, utvalda med slumpfrö 123. Varje block börjar med ett nytt modelltillstånd. Dess första token ger kontext och de återstående 511 utvärderas. Det blir 32 704 utvärderade token per konfiguration. Modell- och datasetrevisioner, tokenisering och programversioner är fixerade i exemplet.⁴

Båda konfigurationerna försämrar modellens förutsägelser kraftigt. Grupp 64 är mindre skadlig här, men inget av resultaten ger oss anledning att bygga en kärna för det här receptet. Packningen har bevarat tritarna exakt. Kvantiseringen har kastat bort för mycket användbar information. Det säger något om vår enkla regel och de vikter vi valde, inte om alla ternära metoder.

Hela körningen, inklusive packning, kontroll av checkpoints och alla tre utvärderingar, tog ungefär 109 sekunder på en Apple M1 Max med modellen och textmaterialet redan i cache. Den tiden gick åt till att besvara kvalitetsfrågan. Den är inte ett mått på inferenshastighet med packade vikter.

Det är jämförelsen som gör arbetet användbart. Du kan variera gruppstorleken, byta regel för skalfaktorn eller ändra vilka tensorer som får receptet och sedan köra samma kvalitetstest igen. Behåll en separat utvärdering med undanhållna data för de konfigurationer du till slut väljer: ett textmaterial som används upprepade gånger för att välja recept har blivit en del av det valet.

Prova ett eget recept

Kodexemplet till artikeln innehåller packning, uppackning, skrivning av checkpoints och utvärdering. På en Mac med Apple Silicon behövs två kommandon för experimentet, efter att du följt installationsanvisningarna och installerat de låsta beroendena med uv sync:

uv run python -m unittest -v
uv run python screen.py --group-sizes 64 128

Skriptet sparar tokenblocken, namnen på de utvalda tensorerna, paketens hashvärden och resultaten under work/. Det anropar MLX LM:s eval_ppl på de sparade blocken med batchstorlek ett. Det håller mätningen liten och undviker att standardvalet av dataset ändras mellan körningar.⁵

Avkodaren tar bara emot paketkatalogen. Den gränsen är värd att bevara när du ersätter exemplet med ditt eget recept.

En befintlig kärna kan räcka

Dina värden kanske ryms i ett format som körmiljön redan implementerar. MLX:s affina kvantisering återskapar ett värde som code * scale + bias.⁶ Koderna $\{0,1,2\}$ med skalfaktorn $s$ och bias $-s$ uttrycker våra ternära värden i dess tvåbitarsformat.

Med en skalfaktor i fp16 och en bias i fp16 per grupp om 64 kostar den nyttolasten $2 + 32/64 = 2{,}5$ bitar per vikt. Den går att köra i det befintliga formatet. Vår layout med fem tritar per byte behöver fortfarande eget stöd. Värden med oregelbundna avstånd passar i allmänhet inte i den här avbildningen. Samma rekonstruerade värden garanterar inte heller identiska flyttalsresultat från olika kärnor.

Vi har inte gjort det experimentet här. Det är ett alternativ efter kvalitetstestet med uppackade vikter, när värdena är värda att gå vidare med.

Slutsats: testa receptet först

Våra ternära paket återskapade varje vald trit exakt, men den rekonstruerade modellens förutsägelser försämrades kraftigt. För det här receptet är nästa steg att se över vilka vikter vi kvantiserar och hur, till exempel genom att behålla tritar i vissa lager och använda MLX:s befintliga format i andra. En snabbare kärna skulle fortfarande räkna med samma skadade vikter.

Det är det beslutet som turen fram och tillbaka hjälper oss att fatta. Börja med en checkpoint som går att köra, ändra bara de avsedda vikterna och jämför modellen som avkodats från det faktiska paketet med dess egen referens. Ett dåligt resultat ger något konkret att ändra innan du skriver en kärna. Ett lovande resultat ger anledning att undersöka vidare.

Perplexitet på den här texten avgör inte kvaliteten på andra uppgifter, språk eller längre kontexter. Testet prövar inte kvantiserade aktiveringar eller den nyckel/värde-cache som används vid generering. En implementation som räknar direkt med formatet behöver också egna kontroller av aritmetik, hastighet och minnesanvändning. De frågorna finns kvar, men vi kan välja vilka recept vi ska lägga det arbetet på.

Fotnoter

1. MLX:s dokumentation av quantize listar de affina bitbredderna. Dess datatypsreferens anger flyttalstypernas bredder. Tabellen är en avgränsad jämförelse, inte en fullständig lista över MLX:s kvantiseringslägen. Byteantalen kontrollerades i MLX 0.32.0 med en bf16-array på 64 gånger 64 element, genom att summera de packade koder, skalfaktorer och bias som quantize returnerar med gruppstorlek 64. Det kontrollerar lagring, inte modellkvalitet eller inferenshastighet. ↩

2. PyTorchs FakeQuantize dokumenterar simuleringen av kvantisering och avkvantisering med flyttal som utdata. ↩

3. Huang med flera beskriver fembitarskodningen i avsnitt 3.1 av Sherry: Hardware-Efficient 1.25-Bit Ternary Quantization via Fine-grained Sparsification. Sherry kombinerar strukturerad gleshet med kvantiseringsmedveten träning och hårdvaruanpassad packning. Vårt enkla exempel med kvantisering efter träning implementerar eller utvärderar inte den metoden. ↩

4. Uppmätt 2026-09-14 på en Apple M1 Max, 32 GiB, macOS 26.5.1, med MLX 0.32.0, MLX LM 0.31.3, NumPy 2.5.1 och Transformers 5.17.0. De 72 utvalda matriserna innehåller 264 241 152 av textmodellens 752 393 024 parametrar. Resultatfilen dokumenterar modell- och datasetrevisioner, hashvärden för indata och källkod, utvalda block och tidsmätning. README anger tokeniseringen. Texten används inte för träning eller kalibrering. ↩

5. MLX LM:s eval_ppl-implementation beräknar korsentropin för nästa token och tar exponentialfunktionen av medelvärdet över token. Vi använder den utvärderaren med egna fasta indatablock. Värdena bör inte jämföras direkt med tabeller som använder annan tokenisering, kontextlängd eller urvalsmetod. ↩

6. MLX:s dokumentation av quantize beskriver affina grupper med skalfaktorer och bias. Bitåtgången här förutsätter att båda lagras med 16 bitar och att tvåbitarskoderna packas utan extra utfyllnad. ↩