Large Parallelism Post: Part III. Mixed Precision Training
В статье предлагается способ снижения потребления памяти тренироки сетей в 2 раза за счет перевода величин из формата FP32 в FP16. Авторы статьи исследуют распределения полученных величин по экспоненте и вводят Loss Scaling для предовращения зануления значений активаций и увеличения точности. Последним дополнением является Arithmetic Precision - разграничение проводимых операций в FP32 и FP16 формате.
Source: Arxive
Современные системы обучения глубокому обучению используют single-precision format FP32. Авторы предлагают тренировать сети используя IEEE half-precision format FP16 - то есть тратить в 2 раза меньше битов, а значит и в 2 раза меньше памяти, сохраняя при этом точность.
FP32 Master Copy of Weights
Сам пайплайн mixed precision довольно прост и изображен на картинке ниже. Во время обучения веса, активации и градиенты хранятся в точности FP16, но для сохранения уровня точности создается Master Copy весов в FP32 формате, которая в последствии обновляется.
Хотя Master Copy весов FP32 далеко не всегда выгодно хранить, для этого есть причины. Первая: во многих сетях значения обновлений (веса умноженные на learning rate) могут быть меньше 2^(-24) - то есть они становятся просто нулем в оптимизаторе FP16. На гистограме ниже видно, что примерно 5% обновлений зануляются в half-precision format, что может негативно повлиять на точность модели.
Вторая: веса могут занулиться в случае большой разницы между их значениями и обновлениями. Даже если значение весов можно представить в формате FP16, но значение обновления будет довольно большим, то при сложении этих параметров результат выйдет за границы формата FP16 и станет нулем, который невозможно будет восстановить. Именно Master Copy весов помогает избежать этого.
Хоть Master Copy весов FP32 увеличивает на 50% потребление памяти по сравнению с ипользованием только FP16, влияние на общее потребление памяти все равно остается куда меньшим, чем тренировка сети только в FP32 формате. Сохранение активаций каждого слоя при Backprop является самой затратной по памяти операцией, поэтому перевод этого процесса в FP16 снижает потребление примерно в 2 раза.
Loss Scaling
Ниже представлена диаграмма доли значений актиаций от экспоненты (2 в соответствующей степени). Легко заметить, что большинство значений выходят за пределы диапазона формата FP16, а значит они зануляются.
Чтобы решить эту проблему, авторы предлагают увеличить экспоненты в 3 раза (масштабировать в 8 раз) - этого будет достаточно для соответствия точности FP32 формата.
Для эффективного сдвига значений в FP16, достаточно масштабировать Loss, вычисленный за Forward pass, перед началом вычисления Backward pass. Chain rule позволяет сохранять масштабирование, а значит не потребуется проводить дополнительные вычисления. Величина масштабирования выбирается эмпирически.
Arithmetic Precision
Для поддержания точности сетей, авторы обнаружили, что в некоторых сетях необходимо, чтобы векторное произведение накапливалось в формате FP32, а только потом переводилось в FP16 перед записью в память. Скорее всего это вызвано ограничениями точности арифместических операций в FP16. Тензорные ядра в архитектуре GPU Nvidia позволяют накапливать dot-product либо в FP16, либо в FP32 формате - Nvidia tesla v100 gpu architecture.
Большие суммы по элементам вектора (такие как в batch-normalization или softmax) также должны выполняться в FP32 формате. То есть чтение и запись в память производится в FP16, а арифметика в FP32 для сохранения точности. Однако, это не замедляет скорость вычислений, поскольку такие операции ограничены пропускной способностью памяти.
Results
Эксперименты проводились для двух подходов:
- Baseline (FP32) - активации, веса и градиенты хранятся в single-precision формате. Все вычисления также проводятся в FP32.
- Mixed Precision (MP) - хранение в памяти и некоторые вычисления осущетсвляются в FP16. Веса, активации и градиенты хранятся в FP16, используется Master Copy весов в FP32. Для некоторых сетей используется Loss Scaling. Используются операции Tensor Core с накоплением в FP32 для сверток, полносвязных слоев и матричных умножений в рекуррентных слоях.
По таблице сравнения точности видно, что MP нередко немного превосходит Baseline вычисления.
Обратите внимание, как выбор величины Loss Scaling влияет на тренировку модели.
Также эксперименты проводились с моделями CNNs Detection, Speech Recognition, Machine Translation, Language Modeling, GANs - все они показали результаты на уровне Baseline.