Full BF16 трейны и при чём здесь Stochastic Rounding

А помните не очень далёкие времена, когда учить что-то в сырой низкой точности никто даже не пробовал? Заводили Mixed Precision, хранили параметры модели в двух точностях - в FP32 и в FP16, смотрели на потребление видеопамяти и немного грустили, ну а что делать, зато работало.
Сейчас, благо, многие уже ставят свои эксперименты в full BF16, - и многие из этих экспериментов даже сходятся из коробки. Но иногда не сходятся. Возникает вопрос, - что же делать в таком случае.

В целом, есть довольно стандартные инженерные трюки, чтобы сделать чуть лучше

* аккумулировать и хранить градиенты в FP32 (и в рамках DP-группы, и в рамках последовательных forward-backward-ов); кстати, ребятишки из PyTorch в следующем релизе 2.10 (во всяком случае это уже есть в nightly сборке 2.10) добавят поддержку раздельного dtype для параметра и для градиента, связанного с этим параметром, - очень ждём 🫶
* кастить состояния оптимизатора, градиенты и параметры в FP32 во время обновления (не эквивалентно тому чтобы хранить их в FP32), - многие реализации оптимизаторов уже так делают, см например оптимизаторы в torchao (ТЫК)

Есть ли какие-то академические работы, которые исследуют феномен проблем, связанных с BF16 трейнами?

Есть! Вот, например, довольно старая, но классная работа (ТЫК), в которой авторы описывают, что проблемы возникают в первую очередь из-за ошибок округления при обновлении параметров модели, а ошибки округления во время расчёта градиентов (forward/backward) имеют минимальное влияние. Авторы предлагают использовать в том числе Stochastic Rounding для того, чтобы исправить эту проблему. Есть ещё одна работа (ТЫК), в которой авторы на практике проверяют full BF16 рецепт со Stochastic Rounding'ом, связывают SR и регуляризацию, а также предлагают реализацию AdamW с SR.

Как работает Stochastic Rounding?

Идея проста и лежит в противовес стандартному округлению. Стандартное округление даёт меньшую ошибку точечно, однако это округление является смещённым. Суть Stochastic Rounding в том, что вероятность округления числа вверх или вниз зависит от того, насколько близко оно к границам сетки точности. В результате математическое ожидание округленного значения равно исходному точному числу (E[SR(x)]=x). Это делает округление несмещенным.

Как реализовать?

Реализация самого SR на pytorch довольно проста. Однако убедитесь, что все ваши процессы в распределённом сетапе имеют одинаковый источник RNG, чтобы обновления были одинаковыми по всему вашему world.

@torch.compile
def copy_fp32_to_bf16_stochastic_(
        target: torch.Tensor, source: torch.Tensor
):
    result = torch.randint(
        0, 1 << 16, source.shape, device=source.device, dtype=torch.int32
    )
    # add the random number to the lower 16 bit of the mantissa
    result.add_(source.view(dtype=torch.int32))
    # mask off the lower 16 bit of the mantissa
    result.bitwise_and_(-65536)  # -65536 = FFFF0000 as a signed int32
    # copy the higher 16 bit into the target tensor
    target.copy_(result.view(dtype=torch.float32))

А полный код оптимизатора, который использует SR, можно подсмотреть в torchao - ТЫК.

В нём авторы кастят параметры, состояния и градиенты в FP32, следовательно делают расчёт обновления в FP32, а затем делают обновление параметров модели уже в BF16, опционально используя Stochastic Rounding.