Den dolda avrundningen i RMSNorm

En typkonvertering i vår RMSNorm-implementation gjorde ingenting. Värdet var redan float32, så att konvertera det till float32 lämnade det oförändrat. Testerna gick igenom.

Med bfloat16 fick samma rad något att göra. Den avrundade ett mellanresultat före multiplikationen med en inlärd vikt. Vår referens behöll mellanresultatet i float32 tills multiplikationen var klar. Formlerna stämde överens, men implementationerna gjorde olika numeriska val.

Det är en användbar detalj när man flyttar en modell mellan plattformar eller utvärderar en kvantisering. Innan man tillskriver de kvantiserade vikterna en avvikelse är det värt att kontrollera vad operationerna runt omkring gör. Vi följde den här skillnaden genom ett litet experiment och källkoden till flera implementationer.

En välbekant operation

RMSNorm (root mean square normalization) håller skalan på en aktiveringsvektor under kontroll. Den finns i många språkmodeller, bland annat Llama och Gemma. För en vektor med $n$ element delar den först varje värde med roten ur medelvärdet av elementens kvadrater. Ett litet $\epsilon$ läggs till under roten för numerisk stabilitet:

$$\hat{x}_i = \frac{x_i}{\sqrt{\frac{1}{n}\sum_{j=1}^{n} x_j^2 + \epsilon}}$$

Sedan skalas varje element med en inlärd vikt, ofta kallad gain:

$$y_i = \hat{x}_i \cdot \gamma_i$$

Formeln lämnar en praktisk fråga obesvarad: hur mycket precision behåller vi längs vägen?

Vårt exempel använder bfloat16, eller bf16, ett 16-bitars flyttalsformat som ofta används för modellvikter och aktiveringar. Det täcker ungefär samma storleksordningar som float32, men har färre bitar för att skilja närliggande värden åt. Att flytta ett float32-resultat till bf16 kan därför avrunda det, även när talet är långt från formatets största möjliga värde.

Två ställen att avrunda

Anta att normaliseringen har gett $\hat{x}_i$ i float32 och att vikten $\gamma_i$ redan är lagrad i bf16. Då finns två möjliga konventioner.

Round once: multiplicera det normaliserade värdet med vikten i float32 och avrunda sedan resultatet till bf16 vid lagring.

$$y_i = \operatorname{bf16}(\hat{x}_i \cdot \gamma_i)$$

Round twice: avrunda först det normaliserade värdet till bf16, multiplicera sedan med vikten och avrunda produkten till bf16.

$$y_i = \operatorname{bf16}(\operatorname{bf16}(\hat{x}_i) \cdot \gamma_i)$$

Här räknar ”once” och ”twice” avrundningarna till bf16 efter normaliseringen. Även float32-beräkningarna har sina avrundningar. Uttrycken beskriver var precision går förlorad, utan att ange hur processorn utför själva multiplikationen.

Det extra parentesparet är lätt att missa vid en kodgranskning. Med float32 som in- och utdatatyp gör konverteringen av mellanresultatet ingenting. Ett test med enbart float32 kan gå igenom utan att alls pröva det valet.

En skillnad du kan återskapa

Vi kan isolera valet utan att köra en modell. NumPy-funktionen nedan avrundar ändliga float32-värden till bf16-precision med round-to-nearest-even. De avrundade värdena ligger kvar i float32-arrayer så att vi kan undersöka dem. Funktionen räcker för det här exemplet, men är ingen fullständig bf16-implementation för specialvärden som NaN.

import numpy as np

def bf16(value):
    # Round finite float32 values to bf16 precision.
    value = np.asarray(value, dtype=np.float32)
    bits = value.view(np.uint32)
    rounded = bits + np.uint32(0x7FFF) + ((bits >> 16) & np.uint32(1))
    return (rounded & np.uint32(0xFFFF0000)).view(np.float32)

def round_once(x_hat, gamma):
    return bf16(x_hat * bf16(gamma))

def round_twice(x_hat, gamma):
    return bf16(bf16(x_hat) * bf16(gamma))

x_hat = np.float32(1.1)
gamma = bf16(1.05)  # The stored gain is actually 1.046875.

print(round_once(x_hat, gamma))   # 1.1484375
print(round_twice(x_hat, gamma))  # 1.15625

Några exempel visar samma effekt. Viktkolumnen innehåller de faktiska bf16-värden som används i multiplikationen:

Den extra avrundningen av mellanresultatet kan flytta slutresultatet åt båda håll. I de här exemplen är svaren grannar bland de tal som bf16 kan representera.

Hur ofta skiljer sig svaren?

Som en större illustration drog vi två miljoner syntetiska värden för $\hat{x}$ från en standardnormalfördelning och vikter jämnt fördelade mellan 0,9 och 1,1. De får stå för värdena som går in i den sista multiplikationen. De är inte aktiveringar hämtade från en modell. Experimentet isolerar typkonverteringen och beräknar inte den föregående RMS-normaliseringen.

Kör detta efter det första kodblocket:¹

rng = np.random.default_rng(42)
x_hat = rng.standard_normal(2_000_000).astype(np.float32)
gamma = rng.uniform(0.9, 1.1, x_hat.size).astype(np.float32)

once = round_once(x_hat, gamma)
twice = round_twice(x_hat, gamma)
different = np.count_nonzero(once != twice)
nonzero = once != 0
relative = np.abs(once[nonzero] - twice[nonzero]) / np.abs(once[nonzero])

print(f"Different outputs: {different:,} / {x_hat.size:,}")
print(f"Share: {100 * different / x_hat.size:.2f}%")
print(f"Largest relative difference: {100 * relative.max():.5f}%")

Svaren skiljer sig i 495 539 fall, eller 24,78 %. Den största relativa skillnaden i detta stickprov är 0,78125 %, räknat mot round-once-resultatet.

Det är en skillnad mellan två beräkningar, inte en uppmätt försämring av modellkvalitet. Det betyder inte att en fjärdedel av en modells svar ändras eller att en bestämd procentsats fel läggs på i varje lager. De frågorna kräver tester på själva modellen.

Var implementationerna skiljer sig

Vi hittade båda konventionerna i etablerade implementationer. Tabellen jämför specifika kodvägar, med källor och avgränsningar nedan. Den är ingen regel för varje operation som ett ramverk erbjuder.

Transformer Engine-kärnan visar också varför gränsen spelar roll för FP8, ett 8-bitars flyttalsformat. I kodvägen för FP8-utdata med en skalfaktor per tensor sker normalisering, viktmultiplikation och utdatas skalning i float32 före den sista konverteringen till FP8. Inget mellanresultat i bf16 behöver avrundas igen. Observationen gäller just den kärnan, inte varje FP8-metod eller beräkningsbackend.³

Följ konverteringen i MLX

Gränsen syns i de här raderna i MLX:s Metal-kärna. Här är uttrycket som skriver resultatet:

out[i] =
    w[w_stride * i] * static_cast<T>(thread_x[i] * local_inv_mean[0]);

Deklarationerna ovanför ger sammanhanget: thread_x och local_inv_mean innehåller float32-värden, medan vikten w och utdatan out använder T. I bf16-kodvägen är T bf16:

1. Produkten inuti static_cast<T>(...) är float32. Konverteringen avrundar det normaliserade värdet till bf16.
2. Vikten och det konverterade värdet är nu båda bf16. Multiplikationen ger ett bf16-resultat och avrundar därför igen. Att lagra det resultatet ger ingen tredje avrundning.

Kärnans loopvariant upprepar samma mönster, med x[r + i] i stället för thread_x[i]. Även där ger multiplikationen med det inverterade RMS-värdet i float32 en float32-produkt före konverteringen.

Vi har implementerat en valbar variant i Precisits MLX-fork. Så här väljer den mellan de två uttrycken:

out[i] = PRECISE
    ? static_cast<T>(
          static_cast<float>(w[w_stride * i]) * thread_x[i] *
          local_inv_mean[0])
    : w[w_stride * i] * static_cast<T>(thread_x[i] * local_inv_mean[0]);

Med PRECISE aktiverad sker båda multiplikationerna i float32 fram till den sista konverteringen. Med "en avrundning" menar vi här konverteringen till bf16; float32-beräkningar har fortfarande sin egen avrundning. Uttrycket grupperar också produkterna som (gain * x) * inverse_rms, i stället för gain * (x * inverse_rms). Att ta bort mellanresultatet i bf16 garanterar alltså inte identiska mellanresultat i float32.

Vilken konvention hör till modellen?

Hugging Faces implementationer av Llama och Gemma dokumenterar skillnaden. I referensversionerna som länkas nedan konverterar Llama det normaliserade värdet före multiplikationen med vikten. Gemma multiplicerar i float32 och konverterar efteråt. Gemma använder också $1 + \gamma$ som gain, vilket är ett separat val i formeln.⁶

Den tidiga konverteringen i MLX följer alltså Llamas konvention vid den här gränsen. Att flytta den till efter multiplikationen skulle ändra beteendet. Färre avrundningar kan föra en beräkning närmare formeln med reella tal, samtidigt som den hamnar längre från implementationen som modellens sparade vikter ska återskapa.

Skillnaden diskuterades redan 2024 i Transformers-projektet.⁶ Att hitta den i våra egna tester blev en påminnelse om att läsa modellkoden tillsammans med ramverkets grundoperation. Ett välbekant operationsnamn kan rymma ett obekant antagande.

Ge testet samma precision som modellen

Vid portning eller jämförelse av kvantiseringar är ett användbart test att ge båda implementationerna samma redan avrundade indata och vikter. Använd datatypen som modellen ska köras med, håll normaliseringens epsilon oförändrat och granska mellanliggande typkonverteringar. Då skiljer man operationens avrundningsval från skillnader som uppstod när dess indata förbereddes.

Om frågan gäller att återskapa en modell, jämför med dess avsedda referenskonvention. Om frågan gäller numeriskt fel, använd en beräkning med högre precision på samma avrundade indata som facit. De besvarar olika frågor.

Utöver multiplikationernas gruppering som vi såg ovan kan även ordningen i summeringar och beräkningen av den inverterade kvadratroten skilja sig. Samma bf16-avrundningsgräns garanterar därför inte bitidentiska utdata. Bakåtberäkningen behöver också egna kontroller.

Slutsats: lär känna dina siffror

När en modell flyttas mellan plattformar byter den också uppsättning numeriska konventioner. För att förstå en jämförelse behöver vi veta vilken precision varje steg använder, var värden avrundas och vilken referens vi vill återskapa. En liten skillnad i utdata är en ledtråd att följa. Dess betydelse för modellkvaliteten behöver mätas separat.

Huaweis Ascend Transformer Boost (ATB) visar ett sätt att göra sådana val synliga. För en stödd fp16-kodväg i RMSNorm väljer precisionMode mellan float32 och float16 i mellanberäkningarna. Det är en separat inställning från modelType för Llama/Gemma och omfattar ett annat val än bf16-avrundningen vi undersökt här.⁷

I vår MLX-fork kan den senare konverteringen väljas med MLX_RMS_PRECISE=1. Ändringen omfattar alla fyra utdatasuttryck i Metal-kärnornas framåtberäkning, både varianten för en rad och loopvarianten, även för ofullständiga rader. Standardbeteendet är bevarat och bakåtberäkningens kärnor är oförändrade.⁸

Med en version byggd från forken ska variabeln sättas innan processen startas. Valet läses vid den första RMSNorm-beräkningen i Metal och behålls sedan för processen. Inställningen gäller hela processen; den finns inte i den upstream-version vi granskade. Dokumentera den valda konventionen tillsammans med modellreferensen när resultat jämförs.

Vårt float32-test besvarade en snävare fråga än vi hade insett. Fallet med lägre precision gjorde det synligt. Att känna sina siffror handlar om att förstå hur de blev till, lika mycket som vad de säger. Nästa gång de nästan stämmer vet vi ett ställe till att leta på.

Fotnoter

1. Återskapat med NumPy 2.5.3. default_rng(42) drar float64-värden som koden sedan konverterar till float32. Vikten avrundas dessutom till bf16 av hjälpfunktionen. Relativ skillnad beräknas som abs(once - twice) / abs(once), med nollor i nämnaren uteslutna. Koden anger fördelning, dragningsordning och avrundning. Detta är ett syntetiskt räkneexperiment, inte ett hårdvaru- eller modelltest. ↩

2. I vår dokumenterade jämförelse med PyTorch 2.14.0 på CPU gav torch.nn.RMSNorm och torch.rms_norm samma resultat som referensen med multiplikation i float32 och konvertering sist, för de testade bf16-indata (4 096 rader, bredd 512). Det fastställer inte beteendet för alla tensorformer eller CUDA-kodvägar. PyTorchs RMSNorm-dokumentation. ↩

3. NVIDIA Transformer Engine, rmsnorm_fwd_kernels.cuh, commit 5f6105b. Den granskade kärnans compute_t är float32. FP8-grenen med en skalfaktor per tensor skalar före konverteringen till utdatatypen. Källkodsgranskning, inte ett GPU-körresultat här. ↩ ↩²

4. Huawei CANN ops-nn, norm/rms_norm/op_kernel/rms_norm_split_d.h, commit fac671fe. I den granskade bf16-kodvägen i ComputeY multipliceras det normaliserade värdet och vikten i float32 före konverteringen till utdatatypen. Källkodsgranskning, inte ett test på Ascend-hårdvara. ↩

5. MLX, rms_norm.metal, commit 6c0f02a. Metal-kärnan konverterar det normaliserade värdet till T före multiplikationen med vikten. Tabellen gäller bf16 för indata, vikt och utdata. ↩

6. Hugging Face Transformers v4.44.2, commit 17489028: LlamaRMSNorm och GemmaRMSNorm. Jämförelsen förutsätter vikter lagrade i indatas datatyp. Gemmas källkod hänvisar uttryckligen till den tidigare diskussionen om avrundningsskillnaden. ↩ ↩²

7. Huawei ATB, RmsNormParam::NormParam, commit 34c40fd9, och RMSNorms API-dokumentation. HIGH_PRECISION_MODE är standard och använder float32 i mellanberäkningarna; HIGH_PERFORMANCE_MODE använder float16. I den granskade headerfilen stöder precisionsvalet float16-indata, utesluter kvantiserade operationer och kan inte kombineras med de alternativa lägena för modelType eller rstd. Det är en närbesläktad API-design, inte belägg för en valbar bf16-gräns mellan round once och round twice eller bitidentiska resultat mellan beräkningsbackends. Inställningarna är dokumenterade; vi har inte prestandatestat lägena på Ascend-hårdvara. ↩

8. Precisits offentliga MLX-fork, commit 4121382 på precise-norm. Template-parametern PRECISE väljer uttryck för framåtberäkningen. Metal-koden som väljer kärna läser MLX_RMS_PRECISE en gång och väljer motsvarande kärna. Valet gäller RMSNorms framåtberäkning i Metal; det lägger inte till varianter för LayerNorm eller CUDA. ↩