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

В этом блоге мы обсуждаем проблемы числовой стабильности при обучении смешанной точности. Часто большие задания по обучению могут быть приостановлены на несколько дней, чтобы справиться с числовой нестабильностью. Мы представляем Tensor Collection Hook для отслеживания условий градиентов во время обучения и показываем, что числовые нестабильности можно выявить раньше, лучше понимая внутреннее состояние модели. Фактически, понимание внутреннего состояния модели на ранних этапах обучения может сообщить, будет ли модель склонна к нестабильности позже в процессе обучения. В наших экспериментах мы можем определить нестабильность градиента в первые часы обучения, за несколько дней до того, как эта нестабильность повлияет на конвергенцию. Основываясь на этих экспериментах, мы предлагаем ряд предупредительных знаков, на которые следует обратить внимание, и способы устранения числовой нестабильности.

Тренировка смешанной точности

Глубокое обучение продолжает двигаться в сторону более крупных базовых моделей. В настоящее время в НЛП доминируют большие языковые модели, такие как GPT и T5. В то время как в CV контрастивные модели, такие как CLIP, показали, что обобщают лучше, чем традиционные модели с учителем. В частности, встроенные в CLIP изученные тексты означают, что он может выполнять логический вывод с нулевым или небольшим числом выстрелов, выходящий за рамки возможностей прошлых моделей CV. Однако обучение базовым моделям может быть проблемой.

Фундаментальные модели обычно включают в себя глубокие сети преобразователей как для зрения, так и для текста, в общей сложности миллиарды параметров на огромных объемах данных. ChatGPT имеет 175 миллиардов параметров. CLIP обучается на сотнях ТБ изображений. Размер как моделей, так и данных означает, что базовые модели требуют недель или даже месяцев обучения на больших кластерах графических процессоров. Чтобы ускорить обучение и сократить количество необходимых графических процессоров, базовые модели обычно обучаются со смешанной точностью.

Обучение смешанной точности помещает некоторые тренировочные операции в FP16, а не FP32. Операции, выполняемые в FP16, требуют меньше памяти и могут выполняться до 8 раз быстрее, чем FP32 на современных графических процессорах. Несмотря на более низкую точность, большинство моделей, обученных в FP16, не демонстрируют какого-либо измеримого снижения производительности из-за их сильной чрезмерной параметризации.

Ускорение обучения с плавающей запятой более низкой точности впервые началось с внедрения тензорных ядер Nvidia в их архитектуру Volta. Поскольку у моделей глубокого обучения так много параметров, точное значение любого параметра часто не имеет значения. Таким образом, представляя числа в 16 битах вместо 32, мы можем одновременно разместить больше параметров в регистрах Tensor Core, увеличивая параллелизм каждой операции.

Но обучение в FP16 представляет собой проблему. FP16 не может представлять числа больше по абсолютной величине, чем 65 504, или меньше, чем 5,96e-8. Фреймворки глубокого обучения, такие как PyTorch, поставляются со встроенными инструментами для работы с ограничениями FP16 (градиентные масштабаторы и автоматическая смешанная точность). Но даже с этими проверками безопасности крупные обучающие задания часто терпят неудачу из-за того, что параметры или градиенты выходят за пределы используемого диапазона. Некоторые элементы глубокого обучения лучше всего работают в FP32. Например, пакетная нормализация часто требует очень тонких корректировок вблизи пределов FP16. Следовательно, пакетная нормализация может либо страдать числовой нестабильностью, либо не обеспечивать достаточную точность для правильной сходимости модели. Это означает, что модели нельзя просто слепо преобразовать в FP16 и ожидать, что они будут обучаться так же, как FP32.

Решение заключается в том, что фреймворки глубокого обучения используют автоматическую смешанную точность (AMP). Это предварительно определенный список операций, которые считаются безопасными для обучения FP16. AMP преобразует части модели, которые считаются безопасными, сохраняя при этом операции, требующие большей точности, в FP32. Кроме того, рекомендуется масштабировать градиенты во время обучения смешанной точности. Некоторые градиенты в модели будут настолько близки к нулю, что окажутся ниже минимального диапазона FP16. Решение состоит в том, чтобы умножить потери на некоторую величину, вычислить более крупные градиенты, а затем отрегулировать масштабирование обратно, когда оптимизатор применяется для обновления весов модели. Ниже приведен пример типичного цикла обучения с включением AMP в PyTorch.

Градиентный масштабатор умножит потери на переменную величину. Если в градиенте наблюдаются NaN, скейлер будет уменьшать множитель наполовину, пока NaN не исчезнут, а затем постепенно увеличивать множитель каждые 2000 шагов, пока NaN не появляются. Это пытается сохранить градиенты в диапазоне FP16, а также не дает градиентам стать равными нулю.

Случаи нестабильности

Несмотря на все усилия, инструменты, встроенные в PyTorch и TensorFlow, не могут остановить все случаи числовой нестабильности, возникающие в FP16.

В реализации HuggingFace T5 более крупные варианты модели давали бесконечные значения даже после обучения. Исследование показало, что в очень глубоких моделях T5 значения внимания накапливаются по слоям и в конечном итоге выходят за пределы диапазона FP16, что приводит к бесконечным значениям и NaN в слое нормализации. Они решили проблему, зафиксировав бесконечные значения на максимальном значении FP16, и обнаружили, что влияние на вывод было незначительным.

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

Beta1 и Beta2 — это параметры скользящего среднего для каждого момента, обычно устанавливаемые на 0,9 и 0,999 соответственно. Деление на бета-параметры в степени номера шага устраняет первоначальную погрешность в обновлении. На этапе обновления мы добавляем небольшой эпсилон ко второму параметру момента, чтобы избежать деления на ноль. Типичное значение по умолчанию для эпсилон — 1e-8. Но помните, что минимальное значение FP16 составляет 5,96e-8. Это означает, что если второй момент слишком мал, обновление будет делиться на ноль. В PyTorch обновление будет пропускать изменения на этом шаге, чтобы обучение не расходилось. Но проблема остается в том, что, особенно при Beta2=0,999, любой градиент менее 5,96e-8 потенциально может остановить обновление веса для этого параметра на длительный период времени. Кроме того, оптимизатор переходит в нестабильное состояние.

Преимущество ADAM заключается в том, что, используя два момента, мы можем адаптировать скорость обучения для каждого параметра. Для медленных параметров обучения мы ускоряем скорость обучения, а для быстрых параметров обучения мы замедляем скорость обучения. Если градиент рассчитывается как ноль для нескольких шагов, даже небольшое положительное значение может привести к расхождению модели до того, как скорость обучения успеет скорректироваться вниз. PyTorch в настоящее время имеет открытую проблему для автоматического изменения эпсилон на 1e-7 при использовании смешанной точности. Это может помочь предотвратить расхождение, когда градиенты возвращаются к положительным значениям. Но при этом возникает новая проблема: увеличивая эпсилон, когда мы знаем, что градиенты находятся в одном диапазоне, мы эффективно уменьшаем способность оптимизатора адаптировать скорость обучения. Увеличение epsilon также не касается случаев, когда обучение останавливается из-за нулевых градиентов.

Градиентное масштабирование в обучении CLIP

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

CLIP одновременно обучает две модели преобразования: языковую модель, подобную GPT, и модель изображения ViT. Глубина обеих моделей создает возможности для градиентов, выходящих за пределы ограничений FP16. И на самом деле реализация OpenClip описывала нестабильность обучения при использовании FP16, когда NaN появлялись в модели около середины обучения.

Крючок для сбора тензора

Чтобы помочь лучше понять внутреннее состояние модели во время обучения, мы разработали Tensor Collection Hook (TCH). TCH оборачивает модель и собирает сводную информацию о весах, градиентах, потерях, входных и выходных данных и состоянии оптимизатора через регулярные промежутки времени.

Например, в этом эксперименте нас интересует состояние градиентов во время обучения. Мы можем захотеть собирать нормы градиента, минимальные, максимальные, абсолютные значения, среднее значение и стандартное отклонение для каждого слоя каждые 10 шагов и визуализировать результаты в TensorBoard.

Затем мы можем запустить TensorBoard с out_dir, установленным в качестве входа --logdir.

Настройка эксперимента

Чтобы воспроизвести нестабильность обучения в CLIP, мы начнем с реализации OpenCLIP и подмножества набора данных изображений Laion 5 миллиардов, которое использовалось для обучения OpenCLIP. Мы оборачиваем модель хуком сбора тензоров, который периодически сохраняет состояние градиентов модели, весов, а также моментов оптимизатора, чтобы мы могли наблюдать, что происходит внутри модели при возникновении нестабильности. Все эксперименты мы проводим на инстансах Amazon SageMaker P4d.24xlarge.

Мы начнем с варианта Vit-H-14, у которого авторы OpenCLIP описали проблемы со стабильностью во время обучения. Начиная с предварительно обученной контрольной точки, мы увеличиваем скорость обучения до 1-e4, аналогично скорости обучения во второй половине обучения CLIP. На 300 шагах в тренировке мы намеренно вводим 10 сложных тренировочных партий подряд.

Потери увеличиваются по мере того, как мы увеличиваем скорость обучения, что и ожидается. Когда сложные случаи вводятся на этапе 300, происходит небольшое, но не существенное увеличение потерь. Модель улавливает сложные случаи, но не обновляет веса на большинстве этих шагов, потому что в градиенте появились значения NaN (показаны треугольниками на втором графике). После прохождения набора сложных случаев градиенты падают до нуля.

Градиентный масштабатор PyTorch

Что здесь происходит? Почему градиенты стремятся к нулю? Проблема заключается в масштабировании градиента PyTorch. Градиентное масштабирование — важный инструмент в обучении смешанной точности. В моделях с миллионами или миллиардами параметров градиент любого параметра будет небольшим и часто ниже минимального диапазона FP16.

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

Чтобы решить проблему потери градиента, ранние методы заключались в том, чтобы просто умножать потери на фиксированную величину, вычислять более крупные градиенты, а затем корректировать обновления весов обратно на те же фиксированные величины (веса все еще хранятся в FP32 во время обучения смешанной точности). Но иногда этой фиксированной суммы все же недостаточно. Более новые методы, такие как средство масштабирования градиента PyTorch, начинаются с большого значения множителя, обычно 65536. Но поскольку оно может быть настолько высоким, что более крупные градиенты могут на самом деле превысить значение FP16, средство масштабирования градиента отслеживает градиенты на наличие NaN, которые указывают на переполнение. Если наблюдаются NaN, обновление веса на этом шаге пропускается, множитель уменьшается вдвое и выполняется следующий шаг. Это продолжается до тех пор, пока в градиентах не перестанут наблюдаться NaN. Если скейлер не обнаружит NaN за 2000 шагов, он попытается удвоить множитель.

В приведенном выше случае масштабатор градиента работает точно так, как предполагалось. Мы передаем ему набор случаев, которые генерируют большие, чем ожидалось, потери, которые создают большие градиенты, вызывающие переполнение. Но проблема в том, что теперь множитель настолько низок, что меньшие градиенты падают до нуля. Скейлер не отслеживает нулевые градиенты, только NaN.

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

Модели, которые производят градиентный нижний поток

Чтобы дополнительно изучить, когда проблема возникает, а когда нет, мы сравнили CLIP с меньшей моделью CV, обычно обучаемой со смешанной точностью, YOLOV5. В обоих случаях мы отслеживали частоту нулевых градиентов в каждом слое в процессе обучения.

В течение первых 9000 шагов обучения от 5 до 20% слоев в CLIP демонстрируют недостаточный градиент градиента, в то время как слои в Yolo показывают только случайный недолив. Скорость недостаточного потока в CLIP также увеличивается со временем, что делает обучение менее стабильным.

Использование масштабатора градиента не решает эту проблему, потому что значения градиентов в CLIP намного больше, чем в YOLO. В случае CLIP, когда средство масштабирования градиента смещает большие градиенты почти до максимума FP16, самые маленькие градиенты все еще падают ниже минимума.

Решение градиентных нестабильностей в CLIP

В некоторых случаях настройка параметров масштабатора градиента может помочь предотвратить недополнение. В случае CLIP мы попытались модифицировать средство масштабирования градиента, чтобы оно начиналось с большего множителя, и сократили интервал увеличения.

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

Одно из возможных решений для улучшения масштабирования — сделать его более адаптивным к диапазону параметров. В документе Адаптивное масштабирование потерь для обучения смешанной точности рекомендуется выполнять масштабирование потерь по слоям, а не для полной модели, чтобы предотвратить недополнение. Наши эксперименты показывают, что необходим еще более адаптивный метод. Поскольку градиенты внутри слоев в CLIP по-прежнему охватывают весь диапазон FP16, масштабирование необходимо адаптировать к каждому отдельному параметру, чтобы обеспечить стабильность обучения. Проблема, однако, в том, что такое детальное масштабирование требует больших объемов памяти, что сокращает размер пакета для обучения.

Более новое оборудование для обучения предлагает более эффективное решение. BFloat16 (BF16) — это альтернативный 16-битный тип данных, в котором точность заменяется большим диапазоном. В то время как FP16 обрабатывает от 5,96e-8 до 65 504, BF16 может обрабатывать от 1,17e-38 до 3,39e38, тот же диапазон, что и FP32. Поскольку BF16 имеет еще более низкую точность, чем FP16, некоторые модели также не сходятся. Но в случае больших моделей трансформаторов BF16 не ухудшает сходимость.

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

Сравнивая CLIP в FP16 и BF16, мы видим только случайное недополнение градиента в BF16.

В PyTorch 1.12 и выше включение BF16 — это небольшое изменение в AMP.

Если требуется большая точность, альтернативой является тип данных Tensorfloat32 (TF32). Представленный Nvidia в графическом процессоре Ampere, TF32 представляет собой 19-битную операцию с плавающей запятой, которая добавляет дополнительные биты диапазона BF16, сохраняя при этом точность FP16. В отличие от FP16 и BF16, он предназначен для прямой замены FP32, а не для смешанной точности. Чтобы включить TF32 в PyTorch, добавьте две строки в начало обучения.

Обратите внимание, что до PyTorch 1.11 TF32 был включен по умолчанию на графических процессорах, поддерживающих этот тип данных. Начиная с PyTorch 1.11, его необходимо включать вручную, чтобы поддерживать ожидаемое поведение для разных поколений графических процессоров. TF32 медленнее обучается, чем BF16 и FP16, с половиной теоретического FLOPS, но все же значительно быстрее, чем обучение в FP32.

BF16 и TF32 доступны на экземплярах P4d, P4de, G5, Trn1 и DL1.

Решение проблем до того, как они возникнут

В приведенном выше примере показано, как определить и исправить ограничение в диапазоне FP16. Но эти проблемы часто возникают только позже, в процессе обучения. В начале обучения модели генерируют более высокие потери и поэтому менее чувствительны к выбросам. Как и в случае с обучением OpenCLIP, могут пройти дни, прежде чем возникнут проблемы, что приведет к потере дорогостоящего вычислительного времени. У FP16 и BF16 есть свои преимущества и недостатки. Ограничения диапазона FP16 могут создавать нестабильность и тормозить тренировку. Но BF16 обеспечивает меньшую точность и может также не сходиться. В идеале мы хотим выявить модели на раннем этапе обучения, которые подвержены нестабильности FP16, чтобы мы могли принять обоснованное решение до того, как возникнут нестабильности.

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

Как модель YOLO, обученная на FP16, так и модель CLIP, обученная на BF16, показывают скорость градиентного недолива, как правило, менее 1%, и остаются стационарными во времени.

Модель CLIP, обученная в FP16, показывает уровень недостаточного потока 5–10% в течение первых 1000 шагов обучения и показывает тенденцию к увеличению с течением времени.

Мы обнаружили, что, используя Tensor Collection Hook для отслеживания скорости снижения градиента, мы можем определить тенденции к более высокой нестабильности градиента в течение первых 4–6 часов обучения. Мы рекомендуем переходить на BF16, когда наблюдаются такие тенденции.

Заключение

Обучение смешанной точности является важным компонентом обучения современных крупных базовых моделей, но требует дополнительного внимания к числовой стабильности. В то время как традиционное обучение модели отслеживает потери и выходные данные, наши эксперименты показывают, что знание внутреннего состояния модели важно для диагностики того, когда модель сталкивается с ограничениями типов данных смешанной точности. Оборачивая нашу модель хуком сбора тензоров, мы можем отслеживать, приближаются ли параметры или градиенты к числовым пределам, и вносить изменения в обучение в ответ или даже до возникновения нестабильности, потенциально сокращая дни неудачных тренировочных запусков. Чтобы начать использовать наш ловушку для сбора тензоров или повторить наши эксперименты, ознакомьтесь с нашей модифицированной реализацией OpenCLIP.