24.3 Mixed precision és nagyobb batch

24.3.1 Cél, előismeret és induló kérdés

Ez az egység a numerikus pontosság és a memóriahatár együttes kezelését vizsgálja. Előfeltétele 24.2. CPU/GPU profiling és helyes időmérés: a mért időintervallum, a gyorsítós szinkronizáció és a memóriamérés fogalmai. Visszautalunk 23.3. Páros összehasonlítás és csoportszintű intervallum páros összehasonlítási szerkezetére, 20.3. Preprocessing, domain shift és súlylicenc preprocessing-szerződésére és 12.3. Sanity check és kis mintára túlillesztés sanity checkjére. Induló kérdés: ha egy $n=16$ elemű batch-et $A=4$ darab $m=4$ elemű mikro-batch-re bontunk, és a gradienseket felhalmozva frissítünk, ugyanazt a paraméterlépést kapjuk-e, mint a teljes 16 elemű batch-csel?

24.3.2 Definíciók és intuíció

Kevert pontosságú tanítás (mixed precision training). A modell paramétereit float32 master-súlyként tároljuk, de a forward és backward egy részét alacsonyabb pontosságú float16 vagy bfloat16 aritmetikával végezzük. A float32 normált értéktartománya hozzávetőleg $10^{-38}$ és $10^{38}$ között van, 23 mantisszabittel; a float16 legkisebb pozitív normált értéke $2^{-14}\approx 6{,}10\cdot10^{-5}$, legnagyobb értéke $(2-2^{-10})\cdot2^{15}=65504$, 10 mantisszabittel. A bfloat16 kitevőtartománya a float32-éhez hasonló (8 kitevőbit), de csak 7 mantisszabittel.

Autocast. A torch.amp.autocast kontextuskezelő, amely operátoronként választja meg, hogy az adott művelet alacsony vagy teljes pontosságon fusson. A kiválasztott pontosságot a device_type és a dtype argumentum adja meg. A PyTorch 2.14 jelölésrendszerében a torch.autocast korábbi elnevezés; az újabb ajánlott felület a torch.amp.autocast (a pontos API-állapotot lásd a forrásjegyzéknél).

GradScaler. A torch.amp.GradScaler a loss értéket egy $S>0$ skálával megszorozza, mielőtt a backward lefutna (alapértelmezés: kezdő skála $2^{16}=65536$, growth_factor=2, backoff_factor=0.5, growth_interval=2000). Ha a skálázatlan gradiens olyan kicsi, hogy float16-ban alulcsordulna, a skálázott $S\cdot\nabla L$ már ábrázolható maradhat. Az unscale_ hívás a gradiensből a skálát visszaosztja, ezután következik az optimizer-lépés. Ha a skálázott gradiensben inf vagy NaN jelenik meg, a scaler.step kihagyja az optimizer-lépést, az update pedig a skálát a backoff_factor-ral csökkenti; ismétlődő túlcsordulásnál a skála 1 alá is mehet, a PyTorch dokumentációja ([D14], „Gradient Scaling” szakasz) kifejezetten jelzi, hogy a skála 1 alá csökkenhet, és ezt a GradScaler nem akadályozza meg. Lab-profilban ellenőrizve: init_scale=1.0 mellett három végtelen gradiensű lépés után a skála $0{,}5$, $0{,}25$, $0{,}125$, a súly pedig változatlan, mert mindhárom lépés kimaradt. Ez heurisztikus dinamikus skálázás, nem általános numerikus garancia. A skála futás közben változó állapot, ezért a tanítás folytathatóságához a scaler.state_dict() is a checkpoint része (12.2. Checkpoint és folytathatóság).

Gradiensakkumuláció. Legyen a teljes batch $D=\{(x_1,y_1),\ldots,(x_n,y_n)\}$, amelyet $A$ darab egyenlő méretű $B_1,\ldots,B_A$ mikro-batch-re bontunk, ahol $|B_a|=m$ és $n=Am$. Az akkumuláció során minden mikro-batch-re kiszámítjuk a mean-redukciójú loss-t, meghívjuk a backward-ot, de a paraméterfrissítést csak $A$ lépés után végezzük el. A mikro-batch loss-t $A$-val osztjuk, vagy a felhalmozott gradienst a frissítés előtt osztjuk $A$-val; a két konvenció nem felcserélhető a clipping és a GradScaler sorrendjében.

Gradiens-clipping. Egy $c>0$ küszöb esetén a teljes paramétervektor $g\in\mathbb R^d$ gradiensére

$$\tilde g = \begin{cases} g, & \|g\|_2 \le c,\\[2pt] c\,\dfrac{g}{\|g\|_2}, & \|g\|_2 > c. \end{cases}$$

A norma a modell összes paraméterének gradiensét egyetlen vektorként kezeli. Kevert pontosságban a clippinget az unscale_ után kell végezni, különben a küszöb $S$-szeres skálán hat.

A hatékony batchméret (effective batch size) az $n=Am$ szorzat; a memóriában azonban egyszerre csak egy $m$ elemű mikro-batch aktivációi élnek. A nagyobb batch a gradiensbecslés varianciáját csökkentheti (6.3. Sztochasztikus és minibatch gradiens), de a memóriaigényt növeli; az akkumuláció ezt a feszültséget oldja fel a frissítési ritmus megtartása mellett.

24.3.3 Tétel: a gradiensakkumuláció ekvivalenciája

Tétel (gradiensakkumuláció ekvivalenciája). Legyen $f(\cdot;\theta)$ a modell paraméterezett függvénye, $\theta\in\mathbb R^d$. Tegyük fel, hogy

(F1) minden $i$-re a per-minta veszteség $\ell_i(\theta)=\ell(f(x_i;\theta),y_i)$ differenciálható $\theta$-ban;

(F2) $f(x_i;\theta)$ értéke nem függ a bemeneti batch többi elemétől, és a modell nem tartalmaz batch-szintű állapotot vagy a hívások közötti véletlen műveletet;

(F3) a mikro-batch-ek azonos $m$ méretűek, $n=Am$, és a paramétereket nem frissítjük a mikro-batch-ek között.

Ekkor a mikro-batch mean-loss gradienseinek átlaga megegyezik a teljes batch mean-loss gradiensével:

$$\frac1A\sum_{a=1}^A \nabla_\theta\!\left(\frac1m\sum_{i\in B_a}\ell_i(\theta)\right) = \nabla_\theta\!\left(\frac1n\sum_{i=1}^n\ell_i(\theta)\right).$$

Bizonyítás (saját levezetés). A gradiens linearitása miatt a jobb oldal

$$\nabla_\theta\!\left(\frac1n\sum_{i=1}^n\ell_i\right) = \frac1n\sum_{i=1}^n\nabla_\theta \ell_i(\theta).$$

Az (F2) feltétel garantálja, hogy minden $\ell_i$ ugyanaz, mint a teljes batchben számolva; az (F1) differenciálhatóság pedig a gradiens létezését. A bal oldal:

$$\frac1A\sum_{a=1}^A \frac1m\sum_{i\in B_a}\nabla_\theta \ell_i(\theta) = \frac1{Am}\sum_{a=1}^A\sum_{i\in B_a}\nabla_\theta \ell_i = \frac1{Am}\sum_{i=1}^n\nabla_\theta \ell_i.$$

Az (F3) miatt $Am=n$, tehát a két oldal egyenlő. $\square$

Megjegyzés (eltérő mikro-batch méretek). Ha a mikro-batch méretek $m_a$ nem egyenlők, a súlyozatlan átlag általában nem egyezik a nagy batch gradiensével. A korrigált súlyozás $\sum_a \frac{m_a}{n}\nabla L_{B_a}$ adja a nagy batch mean-gradiensét, mert

$$\sum_a \frac{m_a}{n}\cdot\frac1{m_a}\sum_{i\in B_a}g_i = \frac1n\sum_i g_i.$$

Az akkumulációs hurkok ezért tipikusan azonos mikro-batch méretet használnak, vagy a loss-t a batch_size / n aránnyal súlyozzák.

Mi töri meg az egyezést? A tétel (F2) feltétele kimondottan szigorú. Két gyakori réteg ellenpélda:

1. BatchNorm train módban. A BatchNorm a forward során a mikro-batch saját átlagát és varianciáját használja. Az $x_1=1$, $x_2=3$, $x_3=5$, $x_4=7$ négyelemű batch átlaga $\mu=4$; a $\{x_1,x_2\}$ mikro-batch átlaga $\mu_1=2$. A normalizált $x_1$ érték a két esetben eltér, ezért a loss és a gradiens is eltér. A BatchNorm tehát nem batch-független a train módban.

2. Dropout. A dropout minden forward híváskor új véletlen maszkot húz. A nagy batch egyetlen forward híváskor $n$ maszkot fogyaszt; az akkumuláció $A$ darab $m$ elemű forward híváskor más RNG-sorrendben fogyaszt maszkokat. Még azonos seed mellett sem garantált, hogy a realizált dropout-maszkok megegyeznek a nagy batch realizációjával. A várható érték ugyanaz lehet, de a realizált paramétertrajektória nem.

A BatchNorm és dropout miatt tehát az azonos effektív batchméret nem garantál azonos tréninget.

24.3.4 Kézi számpélda

Legyen az egyparaméteres modell $f(x)=wx$, a négyzetes per-minta veszteség $\ell_i(w)=(wx_i-y_i)^2$, a batch $D=\{(1,2),(2,4),(3,6),(4,8)\}$, és $w=0{,}5$. Teljes batch mean-loss:

$$L_D(w)=\frac14\sum_{i=1}^4(wx_i-y_i)^2,$$

gradiense

$$\nabla_w L_D = \frac14\sum_i 2(wx_i-y_i)x_i.$$

A négy tag: $2(-1{,}5)\cdot1=-3$; $2(-3)\cdot2=-12$; $2(-4{,}5)\cdot3=-27$; $2(-6)\cdot4=-48$. Összegük $-90$, tehát

$$\nabla_w L_D = \frac{-90}{4} = -22{,}5.$$

Most $A=2$, $m=2$: $B_1=\{(1,2),(2,4)\}$, $B_2=\{(3,6),(4,8)\}$. A mikro-batch gradiensek:

$$\nabla L_{B_1}=\frac{-3-12}{2}=-7{,}5,\qquad \nabla L_{B_2}=\frac{-27-48}{2}=-37{,}5.$$

Az átlaguk $(-7{,}5-37{,}5)/2=-45/2=-22{,}5$, ami pontosan a teljes batch gradiens. A példában nincs BN vagy dropout, ezért a tétel alkalmazható.

24.3.5 Kód: f32 kontroll és jelölt AMP-terv

Az alábbi kód a f32 CPU-főutat ellenőrzi. A kézirat írásakor nem futtattuk; a kimenet a hallgató helyi futtatásának tárgya, az egyezés a tételből következik.

# Futtatás: cd lab && python n095_accum_check.py
import torch
import torch.nn as nn

torch.manual_seed(20261020)
device = torch.device("cpu")
dtype = torch.float32

class TinyMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Flatten(),
            nn.Linear(24 * 24, 16),
            nn.ReLU(),
            nn.Linear(16, 3),
        )

    def forward(self, x):
        return self.layers(x)


def flattened_grad(model):
    return torch.cat([p.grad.reshape(-1).clone()
                      for p in model.parameters()])


model = TinyMLP().to(device)
criterion = nn.CrossEntropyLoss()  # mean reduction, int64 target

x = torch.rand(16, 1, 24, 24, device=device, dtype=dtype)
y = torch.randint(0, 3, (16,), device=device)

# Referencia: teljes n=16 batch
model.zero_grad()
loss_full = criterion(model(x), y)
loss_full.backward()
g_ref = flattened_grad(model)
print("reference grad norm:", g_ref.norm().item())

# Akkumuláció: A=4, m=4, n=16
model.zero_grad()
A = 4
m = 4
for a in range(A):
    xb = x[a * m:(a + 1) * m]
    yb = y[a * m:(a + 1) * m]
    loss_a = criterion(model(xb), yb)
    (loss_a / A).backward()   # a mikro-batch mean loss osztva A-val
g_acc = flattened_grad(model)
print("accumulated grad norm:", g_acc.norm().item())
print("max abs diff:", (g_ref - g_acc).abs().max().item())

A loss_a / A osztás azért szükséges, mert a CrossEntropyLoss mean-redukciójú: minden mikro-batch loss a saját $m$ elemű átlaga. Az $A$ darab ilyen átlag összege $A$-szorosa a kívánt $n$ elemű átlagnak, ezért az $A$-val osztás állítja helyre a nagy batch mean-loss skáláját. Mivel a modellben nincs BN vagy dropout, a várt maximális abszolút eltérés a lebegőpontos kerekítés szintjén van; float32 CPU-n az allclose(atol=1e-6, rtol=1e-5) ellenőrzés észszerű.

A GPU-s AMP-terv jelölt, nem futtatott vázlata, a 2.14-es felület szerinti formában:

# CUDA-környezetben önállóan futtatandó időmérési vázlat.
scaler = torch.amp.GradScaler(device="cuda")
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)

for xb, yb in loader:
    optimizer.zero_grad(set_to_none=True)
    with torch.amp.autocast(device_type="cuda", dtype=torch.float16):
        logits = model(xb)
        loss = criterion(logits, yb)
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)
    scaler.step(optimizer)
    scaler.update()

A sorrend indoklása: az unscale_ után a gradiens float32-ben van, ezért a clipping a valódi $c=5$ küszöböt látja; ha a clippinget az unscale_ előtt hívnánk, a küszöb a skálázott gradiensre vonatkozna. A scaler.step csak akkor frissít, ha a gradiensben nincs inf/NaN; a scaler.update ezután igazítja a skálát. CPU-n bfloat16 autocast esetén a GradScaler indoklása gyengébb, mert a bfloat16 kitevőtartománya a float32-éhez hasonló; a támogatottságot a helyi környezetben ellenőrizni kell.

24.3.6 Gyakori tévedés és ellenpélda

Tévedés: „Azonos effektív batchméret mellett az akkumuláció ugyanazt a tréninget adja, mint a nagy batch.”

Ellenpélda. Tekintsünk egy Linear → BatchNorm → ReLU → Linear modellt. A tétel (F2) feltétele sérül: a BatchNorm train módban a mikro-batch statisztikáját használja. A fenti négyelemű $1,3,5,7$ példában a nagy batch átlaga $4$, az első mikro-batch átlaga $2$; a normalizált bemenetek és így a gradiens is különbözik. Dropout esetén az RNG-hívások száma és sorrendje tér el. A két tréning ezért még azonos seed mellett is szétválhat; az effektív batchméret egyenlősége szükséges, de nem elégséges feltétel.

24.3.7 Feladatok

F1 (fogalmi, nehézség 1). Magyarázza el, miért lehet probléma egy $10^{-6}$ nagyságú gradiens float16-ban, és hogyan viszonyul ehhez a GradScaler. Adjon konkrét $S$ értéket, amellyel a gradiens a normált float16 tartományba kerül.

F2 (levezetési, nehézség 2). Legyen $n=6$, $B_1$ mérete $4$, $B_2$ mérete $2$. Mutassa meg általános $g_i=\nabla \ell_i$ vektorokkal, hogy $\frac12(\nabla L_{B_1}+\nabla L_{B_2})\neq \nabla L_D$ általában, és adja meg a korrigált súlyozást, amely pontosan a nagy batch mean-gradiensét adja.

F3 (implementációs, nehézség 2). Írjon accumulate_gradients(model, x, y, A) függvényt, amely a fenti CPU-példa adatain (16 kép, $A=4$, $m=4$) felhalmozza a gradienst, és adja vissza a referencia-gradienstől való maximális abszolút eltérést. A megoldásnak tartalmaznia kell az importot, seedet, dtype-ot, eszközt és egy allclose ellenőrzést.

F4 (hibakeresési, nehézség 2). Az alábbi kódrészlet „gradiensakkumulációnak” nevezi magát:

for micro_x, micro_y in loader:
    loss = criterion(model(micro_x), micro_y)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

Mi a hiba? Milyen következménye van a paramétertrajektóriára, és hogyan javítaná úgy, hogy a hatékony batchméret $Am$ legyen?

F5 (kísérleti/transzfer, nehézség 3). Tervezzen AMP-kísérletet egy helyi GPU-n vagy bfloat16-képes CPU-n a GlyphLab CNN-en. Adja meg a kérdést, a f32 baseline-t, a változó tényezőt, a választási metrikát, a seed- és környezetkezelést, valamint azt a következtetési határt, amelyet a kísérlet nem léphet át.

24.3.8 Megoldások és indoklás

F1. A float16 legkisebb pozitív normált értéke $6{,}10\cdot10^{-5}$; a $10^{-6}$ gradiens ennél kisebb, tehát a subnormal tartományba kerül, ahol a mantissza vezető bitjei elvesznek, vagy nullára kerekül. A GradScaler a loss-t és ezzel a visszaterjedt gradienst $S$-szeresére nagyítja. $S=2^{16}=65536$ esetén $S\cdot10^{-6}=0{,}065536$, ami a normált float16 tartomány belsejében van; az unscale_ a frissítés előtt visszaoszt. A skálázás nem változtat a gradiens irányán, csak az ábrázolhatóságát javítja.

F2. Jelölje $G_a=\nabla L_{B_a}=\frac1{m_a}\sum_{i\in B_a}g_i$. Ekkor

$$\frac12(G_1+G_2) = \frac12\left(\frac14\sum_{i\in B_1}g_i+\frac12\sum_{i\in B_2}g_i\right) = \frac18\sum_{i\in B_1}g_i+\frac14\sum_{i\in B_2}g_i,$$

míg a nagy batch gradiens

$$\nabla L_D=\frac16\left(\sum_{i\in B_1}g_i+\sum_{i\in B_2}g_i\right) = \frac16\sum_{i\in B_1}g_i+\frac16\sum_{i\in B_2}g_i.$$

A $B_1$ tag együtthatója $1/8$ a $1/6$ ellenében, a $B_2$-é $1/4$ a $1/6$ ellenében; általában tehát nem egyenlők. A korrigált súlyozás $\frac{m_a}{n}$: itt $\frac46 G_1+\frac26 G_2=\frac16(\sum_{B_1}+\sum_{B_2})$, ami pontosan a nagy batch gradiens. Szélsőeset: $m_1=n$, $m_2=0$ esetén a korrigált súly $G_1$-re 1, ami triviálisan helyes.

F3. Egy lehetséges megoldás:

import torch
import torch.nn as nn

torch.manual_seed(20261020)
device = torch.device("cpu")
dtype = torch.float32

class TinyMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Flatten(), nn.Linear(24 * 24, 16),
            nn.ReLU(), nn.Linear(16, 3),
        )
    def forward(self, x):
        return self.layers(x)

def flat_grad(model):
    return torch.cat([p.grad.reshape(-1).clone()
                      for p in model.parameters()])

def accumulate_gradients(model, x, y, A):
    model.zero_grad()
    n = x.shape[0]
    m = n // A
    criterion = nn.CrossEntropyLoss()
    for a in range(A):
        xb = x[a * m:(a + 1) * m]
        yb = y[a * m:(a + 1) * m]
        loss = criterion(model(xb), yb)
        (loss / A).backward()
    return flat_grad(model)

model = TinyMLP().to(device)
x = torch.rand(16, 1, 24, 24, device=device, dtype=dtype)
y = torch.randint(0, 3, (16,), device=device)

g_acc = accumulate_gradients(model, x, y, A=4)
model.zero_grad()
criterion = nn.CrossEntropyLoss()
criterion(model(x), y).backward()
g_ref = flat_grad(model)

diff = (g_ref - g_acc).abs().max().item()
print("max abs diff:", diff)
assert torch.allclose(g_ref, g_acc, atol=1e-6, rtol=1e-5)

Alakellenőrzés: g_ref és g_acc dimenziója a paraméterek számának összege, itt $576\cdot16+16+16\cdot3+3=9472$. A ReLU ugyanazon a bemeneti ponton azonos előjelet ad, ezért a nem-sima pont nem érinti az egyezést.

F4. A hiba: a step() és zero_grad() minden mikro-batch után lefut, tehát ez nem akkumuláció, hanem $A$ darab egymás utáni kis batch-es frissítés. A paramétertrajektória eltér a nagy batch-es trajektóriától; BN-nél és dropoutnál a forward is eltérő statisztikát/maszkot lát, tehát még a gradiens sem a nagy batch gradiense. Javítás: a backward()-ot minden mikro-batch-re meghívjuk, de a step() és zero_grad() csak a ciklus után fut le, és a loss-t $A$-val osztjuk (vagy a felhalmozott gradienst osztjuk a frissítés előtt). Ha a clippinget használunk, azt a teljes, $A$-val osztott gradiensre kell alkalmazni.

F5. Példa protokoll. Kérdés: az AMP megőrzi-e a záró test pontosságot a f32 kontrollhoz képest a GlyphLab CNN-en? Baseline: f32 SGD, ugyanaz a seed (pl. 20261020), tanítási séma, batchméret és early stopping a val loss alapján. Változó tényező: AMP be/ki. Metrika: a legjobb val-checkpoint test accuracy-ja; numerikus kontrollként az első lépés gradiensnorma-eltérése. Split: a rögzített scene-szintű train/val/test; a korrupt tesztet csak a végén értékeljük. Környezet: ugyanaz a gép, PyTorch-verzió rögzítve, GPU esetén szinkronizált mérés 24.2. CPU/GPU profiling és helyes időmérés szerint. Következtetési határ: a GlyphLab kicsi és mesterséges; az eredmény nem általánosítható természetes képi benchmarkra vagy más hardverre, és a seed-szórás (23.3. Páros összehasonlítás és csoportszintű intervallum) miatt egyszeri futás nem bizonyít fölényt.

24.3.9 Önellenőrzés és visszahívás

Ellenőrizze, hogy (1) a clippinget az unscale_ után hívja-e; (2) az akkumulációhoz a mikro-batch loss-t osztja-e $A$-val; (3) a modellben nincs-e BN/dropout, mielőtt az egyezést állítja. Emlékeztető 24.2. CPU/GPU profiling és helyes időmérés alfejezetből: a GPU-időt csak szinkronizáció után mérjük, és a memóriaméréshez a foglalás időpontját is rögzítjük. 23.3. Páros összehasonlítás és csoportszintű intervallum szerint a seed-variáció és a páros teszt külön bizonytalanság; 20.3. Preprocessing, domain shift és súlylicenc szerint a preprocessing-szerződés az AMP-től függetlenül érvényes; 12.3. Sanity check és kis mintára túlillesztés tiny overfitje az akkumulációs hibák gyors szűrője.

24.3.10 Rövid összefoglaló

A kevert pontosságú tanítás float32 master-súlyok mellett alacsonyabb pontosságú forward/backward operátorokat használ. A GradScaler a loss skálázásával akadályozza meg, hogy a float16 gradiens alulcsorduljon; a frissítés előtt a skálát visszaosztjuk. A gradiensakkumuláció azonos méretű mikro-batch-ek és batch-független modell esetén ekvivalens a nagy batch átlaggal: a mikro-batch mean-loss gradienseinek átlaga a teljes mean-loss gradiense. BatchNorm és dropout megtöri ezt az egyezést, mert a forward eredménye a batch összetételétől vagy a hívások közötti RNG-állapottól függ. A clipping az unscaling után a globális gradiennormára hat. A labor f32 kontrollt készít; mért AMP-eredményt csak tényleges futtatás után szabad közölni.

24.3.11 Következő alkalmazás és források

A 24.4. VI. ellenőrző pont: megbízható CNN-training VI. ellenőrző pont a teljes CNN-training megbízhatóságát kéri számon: az itteni akkumulációs és AMP-döntéseket a tanulási görbével, checkpointtal és záró teszttel együtt kell indokolni.

Felhasznált források:

• [D02] PyTorch contributors (2.14). Reproducibility. https://docs.pytorch.org/docs/2.14/notes/randomness.html.
• [D14] PyTorch contributors (2.14). Automatic Mixed Precision. https://docs.pytorch.org/docs/2.14/amp.html.