bfloat16 na CPU: połowa pamięci, czterokrotnie dłużej

Odruch z pracy na GPU brzmi: zmniejsz precyzję, będzie szybciej. Na kartach od Ampere w górę bf16 ma rdzenie tensorowe, a mniejsze wagi to mniej przepychania przez pamięć. Przeniosłem ten odruch na CPU przy zadaniu oceny tekstów dwoma małymi modelami językowymi i dostałem wynik odwrotny do oczekiwanego.

Pomiar

Dwa modele po 1,5 mld parametrów, okno 512 tokenów, 4 wątki, batch 1, te same 10 tekstów w obu przebiegach. Maszyna spokojna: load poniżej 4, żadnych innych pomiarów w tle. Procesor i5-12400F pod WSL2.

Pamięć spadła prawie o połowę, wczytanie skróciło się pięciokrotnie, a samo liczenie wydłużyło się 4,3 raza.

Jakość sprawdziłem osobno, bo „szybciej” i „wolniej” nic nie znaczą, jeśli wynik się zmienia. Rangi ocen z obu przebiegów były identyczne: korelacja Spearmana ρ = 1,0000. bf16 nie popsuł ani jednej kolejności. Zepsuł wyłącznie czas.

Dlaczego

Pierwsze miejsce, gdzie warto zajrzeć, to flagi procesora:

grep -o -w -E 'avx2|avx512f|avx512_bf16|amx_bf16' /proc/cpuinfo | sort -u

U mnie wychodzi samo avx2. Brak avx512_bf16 i amx_bf16 oznacza, że procesor nie ma instrukcji, które mnożą wektory bf16 natywnie. PyTorch nadal przyjmie tensory bf16, ale każdą operację wykona przez konwersję: rozszerzy bf16 do float32, policzy, zaokrągli z powrotem. Mnożenie macierzy w float32 trafia w dobrze zoptymalizowane jądra AVX2, a wariant bf16 płaci za konwersję w każdej operacji i nie korzysta z tych samych ścieżek.

Szybsze wczytanie ma prostsze wyjaśnienie. Wagi tych modeli leżą na dysku w bf16. Dla torch_dtype=torch.bfloat16 to po prostu odczyt. Dla float32 każda waga jest rozszerzana do dwukrotnie większego bufora, stąd 16 sekund i 12 GB w szczycie.

Na procesorach z AMX (Sapphire Rapids i nowsze) albo z avx512_bf16 ta tabela może wyglądać zupełnie inaczej. Dlatego zapisałem ją jako fakt o konkretnej maszynie, a nie o bf16.

Reguła, którą z tego wyciągnąłem

Na tym sprzęcie bf16 na CPU wybieram wyłącznie wtedy, gdy brakuje pamięci. Nigdy dla szybkości. Zapisana jako kod reguła jest krótka:

def cpu_dtype(model_bytes_fp32: int, free_bytes: int) -> torch.dtype:
    # Bez sprzętowego bf16 to wybór pamięci, nie prędkości.
    if model_bytes_fp32 * 1.2 < free_bytes:
        return torch.float32
    return torch.bfloat16

Zapas 1,2 raza to margines na aktywacje i bufory pośrednie; przy oknie 512 i batchu 1 jest z nawiązką. Ważniejsze, że decyzja zależy od pamięci, a typ jest efektem, nie ustawieniem „bo tak szybciej”.

Druga pułapka z tego samego dnia: VRAM to nie RAM

Przy okazji rachunku „ile modeli zmieści się naraz” łatwo zsumować dwie różne rzeczy. nvidia-smi pokazuje pamięć karty, /proc/<pid>/status pokazuje pamięć hosta. Model 8B w 8 bitach zajmował 9,7 GB VRAM, a RSS procesu wynosił 0,57 GB. Dodanie tych liczb w jednym budżecie daje ścianę, której nie ma, a ich pomylenie daje odwrotnie: proces, który wygląda na lekki, a zajmuje prawie całą kartę.

Jak mierzyć, żeby nie oszukać siebie

Kilka rzeczy, które zadecydowały, że tym liczbom ufam:

• Ta sama próbka w obu wariantach. Dziesięć tych samych tekstów, ta sama kolejność.
• Mediana, nie średnia. Pierwszy tekst zwykle płaci za rozgrzanie alokatora i cache, a średnia go przemyca do wyniku.
• Jakość obok czasu. Bez porównania rang nie wiedziałbym, czy wolniejszy wariant nie liczy po prostu czegoś innego.
• Spokojna maszyna. Pomiar wydajności w tle innej pracy mierzy tę drugą pracę.
• Wczytanie osobno od liczenia. Dla jednego tekstu z tej próbki całość od startu procesu to 2,9 + 62,5 s wobec 16,0 + 14,6 s, więc bf16 nadal przegrywa. Ale proporcje zależą od tego, ile pracy przypada na jedno wczytanie: jeśli czas liczenia skaluje się z długością tekstu, to przy tekście dziesięć razy krótszym szacunek daje około 9 s dla bf16 i 17 s dla float32. Tego wariantu nie mierzyłem, to ekstrapolacja.

Ostatni punkt jest praktyczny: w narzędziu wiersza poleceń, które ładuje model, ocenia jeden krótki tekst i kończy działanie, bf16 może wygrać (do sprawdzenia pomiarem, nie wiarą). W usłudze, która ładuje model raz i ocenia tysiące tekstów, przegrywa czterokrotnie. Ten sam typ danych, dwie różne odpowiedzi, a o wyborze decyduje profil pracy, nie nazwa formatu.