Love. Death. Transformers.

18 July 2023 11:43

Максимальное сжатие по вайсману?
Если кто то не знал, то gpu используются при обучении в лучшем случае на 50% от теоретического максимума, для того чтобы стать ближе к нему любят всякие страшные хаки - tensor/pipeline paralel, flash attention и прочие: а чо то у меня не работает.

Оригинальный флеш получал основное ускорение за счет простой идеи: мы используем поблочное умножение для Q,V + хранили экспоненты для вычисленных значений, что позволяло экономить на IO.
При этом еще вычисляют все слои внутри одного cuda kernel что экономит на IO + используют быстрый L2 L3 cache.

Второй flash attention это почти алгоритм из центра наружу, авторы выкинули экспоненты и высчитаывают градиенты на слой по диагональным статистикам + только в конце каждой итерации умножения они скейлят градиенты(из block matmul)
+ вместо хранения сум экспонент они хранят только logsumexp;


красивое....