Методы обучения с низкой точностью

Выпуск нового типа оборудования привел к появлению новых парадигм обучения, которые лучше его используют. В настоящее время это проявляется в обучении с точностью 8 бит с использованием таких пакетов, как TransformersEngine (TE) или torchao (нативный PyTorch FP8).

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

Краткая диаграмма

Ниже представлена краткая диаграмма, показывающая различные битовые точности для каждого решения во время обучения:

TransformersEngine

TransformersEngine — первое решение, направленное на обучение в формате 8-битного плавающего числа. Оно работает за счёт использования заменяемых «на лету» слоёв для определённых слоёв модели, которые используют их FP8-движок для уменьшения количества бит (например, с 32 до 8) без снижения итоговой точности модели.

В частности, Accelerate находит и заменяет следующие слои на версии из TransformersEngine:

• nn.LayerNorm на te.LayerNorm
• nn.Linear на te.Linear

В результате мы получаем модель, в которой большинство слоёв находятся в формате BF16, а некоторые слои — в формате FP8, что позволяет сократить часть потребления памяти.

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

TransformerEngine может принимать множество различных аргументов, настраивающих выполнение вычислений в формате FP8 и их поведение. Полный список аргументов приведён ниже:

• margin: Поле, используемое для масштабирования градиента.
• interval: интервал, определяющий, как часто пересчитывается коэффициент масштабирования.
• fp8_format``: The format to use for the FP8 recipe. Must be one ofHYBRIDorE4M3. (GenerallyHYBRIDfor training,E4M3` for evaluation)
• amax_history_len: Длина истории, используемой для вычисления коэффициента масштабирования
• amax_compute_algo: Алгоритм, используемый для вычисления коэффициента масштабирования. Должен быть одним из max или most_recent.
• override_linear_precision: Указывает, следует ли выполнять GEMM-операции fprop, dgrad и wgrad с повышенной точностью.

Вы можете настроить каждый из этих параметров в составе utils.TERecipeKwargs, чтобы оптимизировать производительность ваших моделей.

Если мы обратим внимание на ранее упомянутую диаграмму, то увидим, что TE просто приводит вычислительные слои к FP8, тогда как всё остальное остаётся в FP32. В результате это приводит к максимальному потреблению памяти, но обеспечивает минимальную потерю точности на этапе обучения.