В чём же считать: fp8, fp32 или fp16

Введение

В жизни каждого DL-инженера рано или поздно появляется задача ужать нейронку. В самом простом случае проблема связана с нехваткой памяти видеокарты. Например, на этапе обучения мы используем изображения небольшого размера, но на инференсе приходится обрабатывать огромные изображения, которые нельзя разбить на части (как в задачах Image2Image или Depth Estimation). Иногда цель более сложная — ускорить время инференса сети. И в такие моменты появляется идея использовать другие типы чисел с меньшей разрядностью. Давайте разберёмся, как это делается и с какими трудностями здесь можно столкнуться.

Что такое floats?

Floats нужны для удобного представления чисел в широком диапазоне: от 0.00034 до 100000342.12. Начнём с простого примера — посмотрим, как выглядит число 33 в целочисленном формате. В двоичной системе оно представляется таким образом:

Формула для расчёта:

С числами с плавающей запятой всё немного сложнее. Они состоят из трёх частей: знака, экспоненты и мантиссы. Формула выглядит таким образом:

Где:

• S — знак числа (1 — для отрицательных, 0 — для положительных);
• M — нормализованная мантисса (начинается с 1);
• E — экспонента со смещением на 127.

Пример: число 33 в формате float

Двоичное представление:

1. Структура:
   • 1 бит: знак (S = 0);
   • 8 бит: экспонента (E = 132);
   • 23 бита: мантисса (M = 0.00001…).
2. Расчёт:

Экспонента:

Мантисса:

Итоговое представление:

Зачем такое сложное представление? Оно позволяет охватывать огромный диапазон чисел. Это особенно полезно в задачах, где значения могут сильно варьироваться. Однако в Deep Learning такой большой диапазон часто избыточен, ведь основные вычисления проходят в пределах от -10 до 10.

Float32 и ниже

Что можно сказать про float32?

Плюсы:

• Всегда и везде работает.

Минусы:

• Занимает слишком много памяти, поэтому обычно стараются уйти в другие типы данных.

Float32 — стандартный формат в большинстве случаев. Для научных расчётов он может быть недостаточен, но в DL его хватает с избытком. Следующий логичный шаг — использовать float16.

Float16

Что можно сказать про float16?

Плюсы:

• Широко поддерживаемый старый тип.

Минусы:

• Динамический диапазон плохо подходит для обучения.

В этом формате 5 бит выделено под экспоненту и 10 бит — под мантиссу. Конвертировать сеть в float16 просто: главное, чтобы видеокарта поддерживала этот тип данных. Если точность сети после конверсии падает, её можно дообучить в float16.

Однако float16 имеет серьёзный недостаток: он плохо справляется с числами, близкими к нулю. Это критично для обучения, так как градиенты часто принимают значения меньше единицы и стремятся к нулю. В таких случаях возникает проблема underflow — когда число представляется нулём. Это снижает точность обучения и требует костылей.

Bfloat16

Что можно сказать про bfloat16?

Плюсы:

• Правильно подобранный динамический диапазон (количество бит в экспоненте) подходит для обучения.

Минусы:

• Не самая широкая поддержка.

Чтобы избежать проблем float16, можно использовать bfloat16 (brain float 16), придуманный Google. В этом формате 8 бит выделено под экспоненту (как и в float32), а на мантиссу остаётся только 7 бит. Теперь мы можем представлять значения меньшего порядка в окрестности нуля, ведь бит для экспоненты стало больше. Это уменьшает проблему underflow, но требует совместимости с железом.

TF32

Что можно сказать про TF32?

Плюсы:

• Не требует усилий для включения.

Минусы:

• Не такой сильный прирост по скорости, не сжимает по памяти.

Ещё один вариант — TF32 (TensorFloat32) от Nvidia. Здесь на экспоненту выделено 8 бит, а на мантиссу — 10 бит. По сути это FP19. Конвертация из float32 в TF32 происходит быстро, так как биты мантиссы просто отбрасываются. В некоторых задачах (например, GAN) точность может немного снижаться, поэтому в PyTorch этот режим включается вручную с помощью torch.allow_tf32 = True.

FP8 и дальше

FP8 — ещё один шаг к уменьшению разрядности. Однако при использовании «в лоб» мы сильно теряем в точности. Применение FP8 оправдано в специфичных задачах, например, в трансформерах. В таких сетях часть вычислений (например, механизм внимания) может быть неинтенсивной арифметически, и большая часть времени уходит на перенос данных между памятью и кешем. Хранение весов в FP8 позволяет существенно экономить память (в Flash Atention 3 легко включается). При этом перед умножением веса можно конвертировать в FP16, чтобы не терять в точности. Хотя такая конвертация может занимать время, выигрыш от экономии памяти на фоне затрат времени на перенос данных становится оправданным.

Важно отметить: под FP8 подразумеваются 2 разных типа — E4M3 и E5M2 (по количеству бит в экспоненте). E5M2 лучше подходит для обучения благодаря большему динамическому диапазону, который нужен градиентам, а E4M3 — для инференса (2 бита в мантиссе — это очень мало для весов).

NF4 — это даже не float

Если нужно сжать сеть ещё сильнее, можно использовать NF4 (Normal Float 4).

Этот формат оптимизирован для нормального распределения весов. Вместо равномерного округления чисел в FP4, NF4 распределяет значения по гистограмме с 16 бинами, что уменьшает потери точности.

В чём секрет? Веса в нейронных сетях обычно распределены нормально. Перед записью числа в контейнер NF4 веса нормализуются, чтобы соответствовать этому распределению, а затем размещаются в одном из 16 диапазонов. Это позволяет сохранить компактное и эффективное представление данных.

Операции с NF4 медленные, так как приходится извлекать значения из списка (по сути это специальное кодирование числа, а не стандартное представление). Но главный выигрыш — экономия памяти, что критически важно для задач с трансформерами.

Заключение

Сжатие нейронок — баланс между производительностью и точностью.

Float16, bfloat16, TF32, FP8 и NF4 — инструменты, которые позволяют оптимизировать вычисления под конкретные задачи. Выбор зависит от аппаратной поддержки и требований к точности в вашей задаче.