Обучение прямо в тритах. Часть первая

Компактно хранить готовую модель — половина дела. Я проверяю следующий шаг: можно ли учить её прямо в таком виде, не держа рядом полную копию весов в обычных числах. Веса — это внутренние настройки, которые меняются, когда модель учится на тексте.

В прошлой записи речь шла о скорости ответов большой модели. Здесь другой опыт: обучение маленькой, примерно с миллионом параметров, на видеокарте A10. Пока это проверка способа, а не результат для большого Тернита.

Вес записан несколькими тритами — цифрами со значениями −1, 0 и +1. Такая запись занимает мало места, но двигаться по ней можно только ступеньками. А полезная поправка часто меньше одной ступеньки. Если её округлять, она пропадёт.

Решение похоже на копилку. Возле каждого веса хранится маленький целочисленный счётчик: поправки накапливаются, и лишь когда набралось достаточно, меняется сам вес. Сила поправок подстраивается по общей статистике для группы из 64 весов. На эту служебную память уходит чуть больше байта на вес. Полная копия весов и полный массив их поправок в обычных числах не нужны. Но промежуточные вычисления остаются обычными: полностью целочисленным это обучение не стало.

На трёх случайных стартах модель училась. Для сравнения я дал обычному Adam — распространённому способу обновления весов — те же начальные значения и учебные порции. Он пока лучше предсказывает текст. Весь измеренный объём выделенной памяти составил около 62 МБ против 73 МБ у Adam. Экономия есть, но обещать по ней многократное сокращение памяти большой модели рано.

Со скоростью сначала было хуже: 900 шагов занимали около 13,5 секунды против 2,6 у Adam. Подробный замер показал, что 61% времени вычислительных ядер видеокарты съедают прямые умножения с упакованными весами. Сам расчёт обновлений занимал только 4%.

Пришлось поменять организацию умножения: небольшой блок весов распаковывается внутри вычислительного ядра и используется сразу для нескольких символов. Время этих умножений упало с 8,67 до 0,41 миллисекунды — примерно в 21 раз. Полной распакованной копии модели по-прежнему нет.

Всё обучение ускорилось скромнее: те же 900 шагов заняли 5,20 секунды. И есть цена: в первом полном опыте ошибка предсказания слегка выросла — с 2,854 до 2,880 бита на символ, где меньше значит лучше. Поэтому ускорение уже измерено, а сохранение качества ещё предстоит подтвердить.

Учиться прямо в компактной записи оказалось возможно. Теперь задача — превратить этот результат в воспроизводимую технологию. Когда она будет собрана, выйдет вторая часть.

Андрей Рощупкин · Тернит

Все записи дневника ↗