С развитием технологий искусственного интеллекта и стремительным ростом мощностей вычислительных систем, модели генерации изображений занимают важное место в сфере компьютерного зрения и творчества. Одним из принципиальных аспектов, влияющих на эффективность и точность таких моделей, является выбор форматов числового представления данных при вычислениях. Среди современных форматов больших и малых разрядностей значительное внимание привлекает BF16, или Brain Floating Point 16-bit — тип числового представления с адаптированной точностью и расширенным динамическим диапазоном. Разберёмся в том, почему BF16 становится всё более популярным в диффузионных трансформерах, какие проблемы он приносит, а также как можно эффективно оптимизировать модели генерации изображений, используя альтернативные решения, ориентированные на платформы с ограниченной поддержкой BF16. В последние годы компании, занимающиеся разработкой продвинутых инструментов генерации изображений, активно интегрируют диффузионные трансформеры.
В их числе – Draw Things, команда инженеров которых создала локальную стековую инфраструктуру для обучения и инференса моделей с параметрами от нескольких миллиардов до десятков миллиардов. Такие масштабные модели, включая SD Medium 3.5 с 2.5 миллиардами параметров и HiDream I1 с 17 миллиардами, предъявляют особые требования к управлению вычислительными ресурсами и точностью вычислений. С ростом глубины и сложности архитектур, таких как FLUX.
1, являющийся вариантом MMDiT, наблюдается тенденция к увеличению масштаба активаций. Это может создавать проблемы с переполнением и снижением точности при работе с форматами с ограниченным диапазоном. В этой ситуации формат BF16 становится привлекательным за счёт большего динамического диапазона по сравнению с FP16. Благодаря расширенному диапазону представления чисел, BF16 способен лучше справляться с масштабированием активаций и поддерживать стабильность обучения и инференса на глубоких уровнях модели. Однако BF16 не лишён собственных сложностей.
В первую очередь это связано с уменьшенной длиной мантиссы, которая уступает FP16 и может негативно сказаться на точности в некоторых вычислительных путях модели. Кроме того, аппаратная поддержка BF16 на популярных в последние годы процессорах Apple Silicon, таких как M1 и M2, оставляет желать лучшего. В macOS 15 появилась программная эмуляция BF16, однако она всё ещё работает примерно на 50% медленнее, чем FP16 на этих устройствах. Для пользователей, ориентированных на мобильные или компактные устройства, в которых скорость и энергопотребление являются критичными параметрами, это создает дополнительные препятствия. В ответ на возникшие вызовы инженеры Draw Things сосредоточили усилия на оптимизации поддержки формата FP16, добиваясь на M1 и M2 производительности, сравнимой с более современными процессорами M3 и M4.
Такой подход позволил расширить возможности запуска моделей на большем количестве устройств, не ограничиваясь при этом использованием BF16, что особенно актуально для пользователей, испытывающих проблемы с поддержкой или точностью BF16. Одной из центральных задач при работе с диффузионными трансформерами является корректная обработка слоя нормализации активаций перед возвращением их в латентное пространство. Этот слой, как правило, позволяет активациям изменять масштаб вне пределов FP16, что вызывает сложности. Чтобы обойти это узкое место, Draw Things предлагает повышать точность вычислений на пути накопления активаций до FP32, при этом не несущую основную нагрузку операций, связанных со сравнительно лёгкими вычислениями — элементарными поэлементными операциями, включая слои нормализации. Развернутое решение было найдено в смешанной точности вычислений с преобразованием режима внутри трансформерного блока.
Адаптивный слой нормализации служит четкой границей, по наступлении которой можно выполнять срез переводов в FP16, сохраняя дальнейшую часть блока именно в этом формате. Такой подход доказал свою эффективность на множестве моделей, например, сериях Wan 2.1, где такие меры достаточно для стабильной работы. Тем не менее, в моделях с большим размером промежуточных слоев MLP присутствует дополнительное усложнение. Проекции в MLP часто расширяют размерность, после чего происходит обратное сужение к скрытому измерению.
При этом возникают риски переполнения активаций в FP16. Для устранения таких проблем применяется консервативное масштабирование значений с коэффициентами порядка одной восьмой или четверти – это позволяет сохранить адекватную точность без возникновения искажений. Важно отметить, что FP16 обладает мантиссой приблизительно в 10 бит и, соответственно, примерно на три бита лучше по точности по сравнению с BF16. Следовательно, масштабированное FP16 значение оказывается зачастую предпочтительнее не масштабированного BF16. Для сохранения стабильности вычислений часть операций с масштабированными активациями, включая оконечные методы GEMM (матрично-векторные умножения) и процессы адаптивной нормализации, выполняются с повышенной точностью FP32.
Такой микс форматов является балансом между производительностью и точностью. Внимательное разделение этапов вычислений и выбор форматов в ключевых точках внутреннего потока позволяет добиться оптимальной работы моделей без необходимости полного перехода на BF16. Еще одной важной областью для обеспечения устойчивости численных операций является масштабирование в операциях внимания. Обычно масштаб присутствует в виде множителя 1/√k, где k — размерность векторов внимания. В реализации системы FlashAttention, на которой основана платформа Draw Things, накопление результатов производится в FP16.
Такая схема может приводить к ошибкам из-за переполнения диапазона. Было обнаружено, что применение масштабирования до операции внимания, а не при ее исполнении внутри ядра, существенно снижает риски переполнения и повышает числовую стабильность. Такая стратегия смешанной точности и точечного масштабирования оказалась эффективной на практике. Для различных моделей команда Draw Things подготовила специфические конфигурации с подробной настройкой коэффициентов масштабирования и определением слоёв, подлежащих масштабированию. Например, в модели FLUX.
1 применяется коэффициент масштабирования 8, воздействующий на конкретные двупоточные блоки внутри сети. Для HiDream используется масштабирование с коэффициентом 4, а в модели Wan 2.1 предмасштабирование применяется только к операциям внимания, обходясь без масштабирования активаций в MLP и прочих блоках. Особая ценность такой методологии в том, что она предоставляет возможность эффективно запускать большие, до 17 миллиардов параметров, модели на устройствах с ограниченной поддержкой BF16. Это расширяет доступ к передовым образам генерации на мобильных и компактных вычислительных системах, обеспечивая стабильность, высокую производительность и качественный результат.
В целом, BF16 выступает как важный инструмент в арсенале современных разработчиков систем ИИ, особенно когда речь идёт о глубоких и масштабных трансформерах. Он ускоряет вычисления и способен удерживать широкий динамический диапазон, что особенно актуально в задачах генерации изображений. Несмотря на определённые ограничения в точности и поддержку на аппаратном уровне, грамотный подход к комбинированию BF16 и FP16, а также использование FP32 в ключевых операциях, позволяет добиваться впечатляющих результатов. Будущее развитие аппаратных платформ, в том числе новых поколений Apple Silicon и других чипов с улучшенной поддержкой смешанной точности, несомненно, снизит пороги для использования BF16. В то же время программные стратегии и тщательно подобранные конфигурации оптимизации, подобные тем, что применяются в Draw Things, позволят максимально раскрыть потенциал существующего железа и сделать технологии генерации изображений доступными и удобными для широкой аудитории.
Таким образом, совмещение преимуществ BF16 с прогрессивными алгоритмами масштабирования и смешанной точности выступает не только ключевым в решении технологических задач, но и открывает перспективы для более эффективной работы моделей генерации изображений на устройствах различного класса и назначения.
