ZeRO

ZeRO (Zero Redundancy Optimizer) to rodzina optymalizacji pamięci treningu. Dzieli stan między procesy Data Parallelism, zamiast utrzymywać wszędzie komplet kopii. Nie jest nową regułą wyznaczania aktualizacji wag: można nim rozdzielić stan używany przez Adam.

Trzy kolejne etapy obejmują coraz więcej danych:

1. ZeRO-1: dzieli stan optymalizatora.
2. ZeRO-2: dodatkowo dzieli zredukowane gradienty.
3. ZeRO-3: dodatkowo dzieli parametry modelu; potrzebne części są zbierane na czas obliczeń i ponownie rozdzielane.

Tak definiuje etapy dokumentacja DeepSpeed. Każdy obejmuje optymalizacje poprzedniego. Offload na CPU lub NVMe jest dodatkową możliwością, nie synonimem ZeRO-3.

Co właściwie znika z pamięci?

W recepturze Mixed Precision Training z Adam analizowanej przez Rajbhandari et al., §3.1 i §5, na parametr przypada 16 bajtów: 2 na wagę FP16, 2 na gradient oraz 12 na kopię wagi FP32 i dwa momenty Adam. Przy $P$ parametrach i $N$ procesach idealny rozmiar przechowywanego stanu na proces wynosi:

Wariant	Bajty na proces
Pełne kopie	$16P$
ZeRO-1	$4P+12P/N$
ZeRO-2	$2P+14P/N$
ZeRO-3	$16P/N$

Własne podstawienie: miliard parametrów i cztery GPU dają odpowiednio 16, 7, 5,5 i 4 GB na GPU, przy dziesiętnym GB. Zwiększ liczbę GPU w demonstracji: w etapach 1 i 2 pozostaje część niedzielona.

Te liczby nie są szczytowym zużyciem VRAM. Nie obejmują aktywacji, buforów komunikacji ani chwilowego zebrania parametrów. Activation Checkpointing ogranicza inną część rachunku — zapamiętane aktywacje.

Stan potrzebny do obliczeń trzeba dostarczyć. W analizie §7 pracy ZeRO etapy 1–2 zachowują wolumen komunikacji bazowego DP, a etap 3 zwiększa go około 1,5 raza. To model komunikacji z konkretnymi założeniami, nie obietnica czasu treningu. Oszczędność pamięci i szybkość trzeba oceniać osobno.

Wykorzystuję treści generowane przez AI jako część mojego codziennego procesu nauki.