Ярослав Булатов, Бен Манн, Дариус Лам

TL; DR; мы сделали поездку Transformer-XL эффективно на 128 графических процессорах в облаке Amazon. Код доступен по адресу https://github.com/cybertronai/transformer-xl.

Обзор

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

К счастью, есть простой рецепт, чтобы ускорить этот процесс:

Шаг 1. Получите хорошую модель одной машины

Шаг 2. Запустите N копий модели на N машинах параллельно, синхронизируя на каждом шаге.

Шаг 3. Решите все оставшиеся технические проблемы.

Мы использовали этот рецепт, чтобы сократить время обучения ImageNet с 2 недель до 18 минут. Вы также можете применить ту же оптимизацию для обучения модели за 2 недели, для чего изначально потребовалось бы 4 года, так что вы можете выбрать масштабирование своего исследования, а не время итерации.

За несколько минут обучение модели изображения «в основном решается», в то время как время обучения языковой модели остается препятствием для инноваций. Наша цель состояла в том, чтобы увидеть, насколько мы можем увеличить пропускную способность обучения, распределяя обучение в облаке.

Для модели с одной машиной мы остановились на официальной реализации Transformer-XL. Эта архитектура позволила добиться нескольких современных результатов в языковом моделировании. Авторы упростили воспроизведение своих результатов, выпустив код и инструкции.

После того, как мы воспроизвели точность в их README, мы повторили производительность, чтобы получить увеличение пропускной способности примерно в 2,6 раза без ущерба для качества.

Затем мы расширили эту модель для обучения на нескольких машинах, используя all-reduce для синхронизации градиентов. Распределенные языковые модели отправляют как минимум в 10 раз больше данных за шаг, чем типичные модели изображений, поэтому мы использовали экземпляры AWS p3dn для их сетевого подключения со скоростью 100 Гбит / с. У этих экземпляров было 32 ГБ памяти графического процессора, поэтому мы смогли повысить эффективность обучения еще на 25% за счет увеличения размера пакета для каждого графического процессора. Конечный результат: версия небольшой модели Transformer-XL с 64 GPU обучается примерно в 44 раза быстрее, чем исходная «медленная» реализация с четырьмя GPU. Увеличивая масштабирование оптимизированной модели Transformer-XL среднего размера, версия со 128 GPU тренирует в 13,2 раза быстрее, чем версия с 8 GPU.

Процедура обучения требовала изменений, чтобы предотвратить численное расхождение при больших размерах пакетов, поэтому мы следовали рецепту, предоставленному Google в их бумаге масштабирования 1024 TPU. Мы открыли исходный код этого оптимизатора в виде отдельного репо.

Технические подробности

Набор данных

В первую очередь мы экспериментировали с Wikitext-103, набором данных моделирования на стандартном языке, состоящим из предварительно токенизированного подмножества Википедии. Обучающий набор содержит примерно 100 миллионов токенов и 267 тысяч уникальных слов в 515 МБ. Он использует модифицированную токенизацию Моисея, при которой 0,4% слов заменяются символами ‹UNK›, чтобы ограничить размер словаря. Кодовая база Transformer-XL загружает весь набор данных в память графического процессора в начале обучения.

Чтобы ускорить процесс обучения, мы предварительно закодировали наши наборы данных и запекли их в образы дисков AWS.

Инфраструктура

Мы автоматизировали многие стандартные задачи, которые связаны с созданием прототипа распределенной обучающей системы, и превратили их в инструмент ncluster. Философия заключалась в том, что любая частая задача должна выполняться одним вызовом из командной строки. Некоторые примеры:

1. Подготовьте несколько экземпляров для тестового запуска по спотовым ценам.

2. Исправьте ошибку и перезапустите эксперимент с несколькими машинами менее чем за 30 секунд.

3. Войдите в систему на самой последней машине, чтобы узнать, что она делает.

4. Синхронизируйте папку на удаленном компьютере с локальной папкой.

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

Сеть

Мы пробовали различные конфигурации сети, чтобы найти настройки, которые позволили бы оптимизировать пропускную способность. Запустив тесты iperf3, мы смогли увидеть пропускную способность 93 Гбит / с между машинами, в то время как производительность NCCL all-reduce упала до примерно 50 Гбит / с. Одна из особенностей настройки сети AWS заключается в том, что каждый сетевой поток ограничен примерно 10 Гбит / с, поэтому мы использовали несколько колец NCCL для включения нескольких потоков. Из-за ошибки в последней версии NCCL нам пришлось вернуться к старой версии, которая позволяла использовать несколько потоков, но не использовала пропускную способность эффективно. В частности, количество переданных байтов выросло в 4 раза по мере увеличения количества графических процессоров, чего нельзя ожидать от алгоритма с постоянной пропускной способностью, такого как ring-allreduce. Однако, когда у нас работал LAMB и наш локальный размер пакета превысил 32 ГБ ОЗУ, наши сетевые требования упали настолько, что даже при чрезмерной нагрузке на связь на каждой машине мы не достигли предела пропускной способности.

В какой-то момент мы заметили, что наши прогоны с наиболее эффективной пропускной способностью были неэффективными, при этом производительность упала на 30% (нижний левый график). Некоторое исследование показало, что работа слишком близко к пределу памяти приводила к тому, что кэширующий распределитель CUDA неоднократно исчерпывал память и приводил к дополнительной синхронизации во время сборки мусора (диаграмма ниже справа). Решение состояло в том, чтобы немного уменьшить использование памяти, уменьшив размер пакета, что привело к исчезновению проблемы.

Слева: время в обратном направлении в мс, справа: torch.cuda.memory_cached ()

Мы обнаружили некоторую изменчивость пропускной способности между запусками. Например, неудачный запуск 64 GPU займет 2,5 минуты, чтобы пройти одну эпоху wt103, в то время как «счастливый» 128 GPU может сделать это за 1 минуту. Мы полагаем, что это происходит из-за того, что группа размещения кластера AWS является приблизительной: AWS пытается разместить все экземпляры в группе размещения в локально подключенной конфигурации или «кирпичике», но когда это невозможно, некоторые экземпляры будут перетекать на соседние блоки.

Смешанная точность

Опора на 16-битную точность позволила нам использовать меньше памяти и повысить скорость обработки за счет использования Volta TensorCores. Мы протестировали две формы тренировки fp16 - чистую тренировку fp16 и тренировку со смешанной точностью. В чистом режиме fp16 все операции выполняются с пониженной точностью. Мы не смогли сопоставить цифры точности в чистом fp16.

При обучении со смешанной точностью мы храним основную копию весов в fp32 и обновляем эту главную копию на каждом шаге оптимизатора, а затем приводим веса к fp16 для прямого прохода. Это позволило нам сохранить точность при экономии места. Масштабирование динамических потерь увеличило объем распространяемой информации о градиенте при сохранении числовой стабильности. Для их реализации мы использовали код из пакетов Apex и Megatron, разработанный NVidia.

В синтетических экспериментах мы обнаружили, что тензорные ядра не были полностью активированы, если их размер не был кратен 8.

Рисунок 1. Время умножения матриц n x n. Оранжевый - fp32, синий - fp16. Ось Y показывает среднее значение в секундах за три прогона. При использовании тензорных чисел умножение матриц fp16 выполняется почти в 10 раз быстрее для больших n. Обратите внимание, что при малых n выигрыш в производительности от использования половинной точности незначителен.

Рисунок 2: Однако, когда n сдвигается на единицу, так что n не кратно 8, fp16 работает одинаково с fp32, т. Е. Никакого увеличения производительности не происходит. а мы просто экономим память.

Изменяя размеры наших моделей, мы можем добиться значительного увеличения скорости с 1,2x до 2 за счет преобразования в режим смешанной точности и увеличения размера партии.

ЯГНЕНОК

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

Напомним, что LAMB умножает скорость обучения для каждого веса слоя на отношение (r) нормы весов (r1) к норме шага Адама (r2).

Как интерпретировать левые диаграммы ниже: ось Y показывает временной шаг количества обработанных токенов, причем первый находится вверху. Ось X - сегменты гистограммы, например, N выборок имели значения от 0 до 1, 1 и 2 и т. Д. Ось Z - это частота гистограммы для значений каждого параметра в сети, например, X% весов слоев попали в От 0 до 1 ведра. Вы можете представить себе это как одну гистограмму значений на каждом временном шаге, сложенных во времени, чтобы мы могли видеть, как гистограмма изменяется с течением времени. Правые диаграммы показывают те же данные, но в виде диаграммы с областями. Более темные цвета означают, что большее количество интервалов гистограммы имело значения в заштрихованной области.

Большинство весов параметров (r1, вверху справа) все время близки к нулю, но некоторые из них растут и распространяются, некоторые становятся довольно большими (максимальное ведро 6,5)

Приведенные выше значения шага Адама вначале очень велики, но быстро стабилизируются до значений от 0 до 0,5.

Вышеупомянутые два эффекта приводят к тому, что коэффициент доверия (r, вверху слева) стимулирует рост некоторых параметров в течение длительного времени, достигая максимального значения клипа, равного 10.

В наших тестах мы обнаружили, что LAMB имеет значение даже на одной машине p3dn.24xlarge с глобальным размером пакета 768, в то время как Адам сразу отклонился. Мы также обнаружили, что можем полностью исключить разминку в нашем графике скорости обучения, поскольку в начале обучения градиенты большие, что приводит к небольшому коэффициенту доверия LAMB и снижает скорость обучения для наиболее нестабильных слоев.

Кодирование пар байтов против адаптивного встраивания

Wikitext-103 имеет размер словаря 267 КБ, что требует много места на модели для представления. Мы попробовали переключиться на кодирование пар байтов (BPE), которое показало преимущества в производительности при машинном переводе. Мы использовали реализацию OpenAI’s BPT, поскольку она может обрабатывать последовательности вне словарного запаса.

Transformer-XL имеет очень большой размер словаря (267 КБ), поэтому большинство параметров находится во встроенных токенах. Мы смогли уменьшить размер модели с помощью BPE с размером словаря 50 КБ, что уменьшило нашу конфигурацию небольшой модели с 186 до 75 млн параметров с той же репрезентативной емкостью. Исходный код Transformer-XL использует адаптивный ввод и softmax для уменьшения размера встраивания.

Одна немного хитрая деталь заключается в том, что BPE производила токенизации примерно на 15% дольше, чем Wikitext-103. Это ожидается, потому что Wikitext-103 уже токенизирован (например, разделение «не делайте» на «не делайте». Если бы это было не так, BPE мог бы обрабатывать «не» как отдельный токен, в то время как токенизатор Wikitext-103 В BPE слова вне словарного запаса разбиваются на части слов, которые есть в словаре, даже в Unicode, путем разделения символов на отдельные байты. Такое поведение является источником 15% увеличения длины токенизации.

В наших исследованиях абляции мы обнаружили, что производительность, измеренная в потерях на обучение на токен, была почти одинаковой с учетом длины токенизации. Количество токенов в секунду увеличилось на ~ 10% из-за меньшего размера модели, а использование графического процессора значительно упало для того же размера пакета с ~ 90% до ~ 55%. Это означало, что при масштабировании до более крупных моделей мы могли сохранить больший размер пакета для каждого графического процессора.

Для более честной борьбы мы также сравнили BPE с адаптивными встраиваниями на нашей большой конфигурации модели:

Преобразование для затруднения проверки: math.exp (math.log (26,25) * 1950/1699) = 42,5, так что это впечатляюще эквивалентно! И это несмотря на то, что кодировка BPE, которую мы использовали, была обучена на общем образце сети (набор данных OpenAI Webtext), а не конкретно на Wikitext-103.

Эти две техники совместимы, поэтому было бы интересно их комбинировать. Тем не менее, поскольку BPE настолько гибкий, а код проще, я склонен использовать только его в будущих экспериментах.

Поиск скорости обучения

Узнав об этой идее в fast.ai, часть 2, Бен был рад увидеть, как она работает. С Адамом все получилось хорошо:

Минимум приведенного выше графика соответствует скорости обучения, использованной в статье (0,00025)! И нам нужно было запустить его всего на несколько минут, чтобы обнаружить это, что намного дешевле, чем выполнять традиционный полный поиск по гиперпараметрам, в котором вы бы работали хотя бы в течение эпохи с одним расписанием LR, которое вы собираетесь использовать позже.

Применительно к LAMB искатель значительно переоценил скорость обучения.

Скорость обучения при минимальных потерях серой кривой составила 0,0263; оранжевый (с использованием спада веса 1e-4) был 0,0435. Но когда мы запустили его даже со скоростью обучения 0,01, он все равно разошелся. Нам пришлось полностью опуститься до 0,005, чтобы получить стабильную сходимость. Мы думаем, что это связано с тем, что LAMB приводит к тому, что фактическая скорость обучения по уровням сильно отличается в разное время обучения, в результате чего средство поиска скорости обучения не распространяется. Вероятно, мы могли бы обойти это, если бы дольше использовали поисковик скорости обучения или добавили контрольную точку в процессе обучения. Вместо этого мы решили выполнить более традиционный поиск по гиперпараметрам скорости обучения.

Масштабирование партии

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

Одним из объяснений увеличения начальной скорости сходимости является то, что мы использовали постоянную скорость обучения для разных размеров пакетов, что привело к более высокой скорости обучения на каждый пример. Разрыв между синей и серой линиями в конце составляет около 100 миллионов токенов (1 час), так что это помогло примерно на 10%, даже несмотря на то, что мы не использовали всю емкость графического процессора.

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

Вся Википедия

На графиках ниже показаны потери при обучении и потери при проверке для конфигурации большой модели в Wikitext-103. На приведенной ниже диаграмме видно, что всего за 800 млн токенов (1,5 часа) наши потери при проверке перестали уменьшаться, что говорит о том, что мы переоснащаем набор данных.

По этой причине мы были рады попробовать обучение на всех дампах Википедии NVIDIA, которые примерно в 25 раз больше, чем Wikitext-103.

  • Это 12 ГБ текста, поэтому мы не можем загрузить все это в память графического процессора.
  • Он хранится в формате JSON, поэтому сначала его нужно выровнять, сохранив границы статьи.
  • Он содержит произвольные уникальные токены, поэтому необходимы некоторые уловки, чтобы ограничить размер словаря.

К счастью, Transformer-XL уже содержал многофайловый загрузчик данных для корпуса объемом 1 миллиард слов, поэтому потребовалось лишь немного потрудиться, чтобы он заработал. Обратите внимание: поскольку токенизация различается, сравнивать показатели потерь не следует. Ниже мы показываем два прогона с использованием 64 графических процессоров во всей Википедии со скоростью ~ 600 миллионов токенов в час. Мы обнаружили, что потеря валидации очень чувствительна к графику скорости обучения.

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

Благодарности

работа была совместными усилиями при участии

Бен - конвейеры данных, LAMB, поиск и настройка параметров

Darius - двукратное ускорение на моделях с одной машиной с использованием тренировки смешанной точности

Ярослав - экспериментальная инфраструктура, распределенное обучение

Дополнительная благодарность AWS за предоставление вычислительных ресурсов для этой работы.