Как Amazon Search M5 сэкономил 30 % на стоимости обучения LLM, используя AWS Trainium

Переиздано Платоном

Читают: 0

На протяжении десятилетий Amazon была пионером и инноватором машинного обучения (ML), предоставляя своим клиентам восхитительный опыт. С самого начала Amazon использовала машинное обучение для различных целей, таких как рекомендации книг, поиск и обнаружение мошенничества. Как и в остальной отрасли, развитие ускоренного оборудования позволило командам Amazon разрабатывать модельные архитектуры с использованием нейронных сетей и глубокого обучения (DL).

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

Как и во многих организациях ML, ускорители в основном используются для ускорения обучения и вывода DL. Когда AWS запустила специализированные ускорители в первой версии Вывод AWS в 2020 году команда М5 быстро начала используйте их для более эффективного развертывания производственных рабочих нагрузок, что позволяет сэкономить средства и сократить задержку. В прошлом году AWS запустила AWS Трениум ускорители, которые оптимизируют соотношение производительности и затрат при разработке и создании моделей DL следующего поколения. В этом посте мы обсудим, как M5 удалось снизить затраты на обучение своих моделей на 30 %, и поделимся некоторыми лучшими практиками, которые мы узнали на этом пути.

Экземпляры трейниума

Благодаря достижениям в области специализированных ускорителей Amazon также предлагает привлекательные ускорители в виде AWS Inferentia и Trainium. Как следует из их названий, эти чипы оптимизированы так, чтобы превосходить потребности в рабочих нагрузках логического вывода и обучения соответственно. Для крупномасштабного обучения моделей фундаментов, размеры которых достигают миллиардов параметров, Trainium Экземпляры Trn1 и Trn1n являются идеальным выбором благодаря своим характеристикам. Экземпляры Trn1 оснащены самой современной НейронКор-v2и иметь большой объем вычислительных ресурсов и памяти ускорителя. Экземпляры Trn1n также можно выбрать с учетом большей пропускной способности сети (1,600 Гбит/с), поэтому они идеально подходят для обучения производительности с учетом оптимизации затрат.

Чтобы использовать ускорители, вам нужен программный уровень для их поддержки. С чипами Trn и Inf Нейрон SDK для AWS разблокирует специально созданные ускорители Amazon с помощью PyTorch XLA. PyTorch XLA преобразует нетерпеливый режим PyTorch в реализацию на основе графов ленивого режима. Эти графики затем используются и далее компилируются для использования с ускорителем. PyTorch Neuron (часть Neuron SDK) позволяет пользователям PyTorch обучать свои модели на Trainium NeuronCores с помощью нескольких строк кода.

Модель и рабочая нагрузка

Команда M5 обучает и внедряет базовые модели и универсальные представления, чтобы помочь различным командам Amazon доставлять удовольствие Amazon.com клиенты. Одной из таких моделей является модель кодировщика текста, за которой следует многослойный перцептрон (MLP) с явными или неявными взаимодействиями функций, определяемыми архитектурой нейронной сети с сотнями миллионов обучаемых параметров. Эта модель обучена на миллиардах токенов и используется для генерации миллионов вложений в автономном пакетном режиме вывода. Эти внедрения являются входными данными для сервиса Amazon первого уровня, ориентированного на клиента.

Инфраструктура для производственного трубопровода использует Пакет AWS стратегии справедливого распределения очередей, используя многоузловой кластер trn1.32xlarge с поддержкой EFA в качестве вычислительной машины для обучения модели. Функционально производственный конвейер выполняет постепенное обучение модели, оценку обученной модели и автономный пакетный анализ обученной модели, используя PyTorch в качестве базовой библиотеки DL.

Цели

Удовлетворение наших клиентов является главным принципом. Учитывая характер конвейера, ориентированного на клиентов, крайне важно, чтобы все соглашения об уровне обслуживания (SLA) соблюдались без регрессов. Мы определили два важнейших критерия приемлемости для адаптации нашего существующего производственного конвейера графических процессоров и перевода его на Trainium:

Качество модели – Качество наших моделей напрямую влияет на качество обслуживания клиентов. Мы требуем, чтобы разница в качестве модели между графическим процессором и Trainium составляла менее 0.1%.
Производительность обучения – Мы периодически обучаем наши модели, чтобы предоставить нашим клиентам новейший опыт. Мы требуем, чтобы сходимость моделей была достигнута в течение заранее определенного периода времени (например, 1 недели) для соблюдения наших производственных соглашений об уровне обслуживания.

В следующих разделах мы поделимся своим опытом работы в обратном направлении от этих критериев и нашими знаниями по поддержке производственных рабочих нагрузок в масштабе Amazon.

Учебный скрипт

Прежде чем начать обучение модели, нам необходимо внести изменения в сценарий обучения, чтобы сделать его совместимым с XLA. Учитывая размер модели, мы используем параллельную обработку данных (DDP) для обучения модели. DDP позволяет нам увеличить производительность обучения модели за счет увеличения количества машин, используемых для обучения модели, без каких-либо изменений кода. Мы следовали инструкциям, изложенным в Учебное пособие Neuron PyTorch MLP для добавления конструкций, специфичных для XLA, в наши сценарии обучения. Эти изменения кода легко реализовать. Ниже приведены некоторые важные технические выводы, полученные в результате этого упражнения, которые значительно повысили производительность нашей модели:

Размещение xm.mark_step() – xm.mark_step() компилирует и запускает лениво собранные графы вычислений. Вызов mark_step слишком много раз приведет к большему количеству маленьких графов, тогда как слишком малое количество вызовов приведет к небольшому количеству, но больших графов. В зависимости от вашего приложения производительность и реализация обучения вашей модели будут варьироваться в зависимости от вашего размещения. xm.mark_step(). Наша реализация ставит одно xm.mark_step() после прямого и обратного прохода и один после шага оптимизатора.
Обертывание загрузчика данных с помощью загрузчика многопроцессорных устройств XLA – Это важный шаг, который можно легко пропустить. Загрузчик многопроцессорных устройств torch_xla.distributed.parallel_loader.MpDeviceLoader загружает данные обучения на каждое устройство XLA с возможностью предварительной загрузки и перекрытия загрузки данных при запусках устройства для повышения пропускной способности. Загрузчик устройства также вызывает xm.mark_step() и поэтому может строить графики загрузки данных на устройство с хоста.

Сборник для Трейниума

Традиционно цикл разработки модели с использованием графических процессоров включает внесение изменений в модель или сценарий обучения и непосредственный запуск их на устройстве с графическим процессором. Ускорители, такие как Trainium, использующие XLA, требуют дополнительного шага, прежде чем обучение модели можно будет запустить на ускорителе. Графики вычислений XLA можно запускать только после их компиляции. Как правило, существует два способа выполнения такой компиляции: с опережением времени (AOT), когда вы сначала отслеживаете и компилируете все графики, а затем запускаете их, или Just In Time (JIT), где графики отслеживаются, компилируются и запускаются по мере их поступления. встречаются. Neuron SDK предоставляет обе эти функции «из коробки». Обычно сначала выполняется компиляция AOT. После этой компиляции запускаются графики. Если встречаются новые графики, среда выполнения Neuron перед их запуском вызывает JIT-компиляцию. Для выполнения AOT-компиляции Neuron SDK предоставляет Neuron_parallel_compile, утилита компиляции, которая извлекает графики из пробного запуска обучающего сценария и выполняет параллельную компиляцию AOT.

Важным аспектом компиляции AOT является обеспечение того, чтобы в ходе обучения не создавались новые графы вычислений. Одним из источников новых графов вычислений (и, следовательно, перекомпиляций) являются динамические формы обучающих пакетов во время обучения модели. Мы обнаружили, что использование статических форм и пакетов фиксированного размера исключает необходимость компиляции времени обучения и значительно повышает производительность обучения без какого-либо влияния на точность модели. Введя такие ограничения на обучение, мы заметили, что для отслеживания всех графиков во время компиляции AOT требуется всего 4–5 шагов обучения модели, один шаг проверки модели и однократная проверка модели. Важно отметить, что Neuron SDK постоянно развивается и в будущем будет поддерживать и динамические формы.

Кроме того, скомпилированные графики сохраняются в Постоянный кэш нейронов на диске или в Простой сервис хранения Amazon (Amazon S3) ведро. Это особенно полезно для производственных рабочих нагрузок, где архитектура модели и конфигурация обучения не меняются. Таким образом, накладные расходы на компиляцию возникают только один раз. Использовать кеш так же просто, как установить флаг среды:

export NEURON_COMPILE_CACHE_URL="s3://BUCKET/KEY"

Компилятор Neuron также предоставляет три параметры оптимизации на уровне компилятора (O1, O2, O3), чтобы сбалансировать время компиляции и производительность выполнения модели. O1 обеспечивает оптимизацию ядра вычислительного графа и минимизирует время компиляции, O3 обеспечивает улучшенную пропускную способность выполнения модели за счет увеличения времени компиляции, а O2 (вариант по умолчанию) представляет собой баланс между ними. В нашем случае мы использовали оптимизацию O1 и наблюдали сокращение времени компиляции на 86 % без изменения показателей точности модели, одновременно наблюдая снижение пропускной способности примерно на 5–7 % по сравнению с оптимизацией по умолчанию (O2). В зависимости от варианта использования вы можете выбрать разные уровни оптимизации.

Подводя итог, мы использовали для компиляции следующие флаги:

NEURON_CC_FLAGS="--target trn1 --auto-cast all --auto-cast-type bf16 --model-type transformer --optlevel O1"

Совместимость контрольных точек

Когда компиляция успешно завершится, мы сможем приступить к обучению наших моделей на Trainium. Как упоминалось ранее, мы обучаем наши модели постепенно, то есть загружаем ранее обученную контрольную точку модели и продолжаем обучение с новыми данными. PyTorch и PyTorch XLA обеспечивают плавный переход между ускорителями благодаря совместимости контрольных точек. Гибкость переключения между графическим процессором и Trainium позволила нам легко загрузить предыдущую модель графического процессора и провести обучение на машинах Trainium. Это было важно для того, чтобы мы могли инициализировать нашу модель с использованием лучшей ранее обученной модели без каких-либо простоев производства или потери точности модели.

Поскольку модель графического процессора была сохранена с помощью стандартных утилит сохранения модели PyTorch, мы смогли использовать утилиту загрузки контрольной точки PyTorch для загрузки модели графического процессора на устройства Trainium.

Например, на GPU/CPU вы можете сохранить модель с помощью следующего кода:

torch.save(model.state_dict(), PATH)

Затем вы загружаете модель обратно в Trainium:

import torch_xla.core.xla_model as xm
xla_device = xm.xla_device()
model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(xla_device)

Аналогичным образом вы можете сохранить модель в Trainium с помощью следующего кода:

import torch_xla.core.xla_model as xm
# automatically moves the data to CPU for the master device
xm.save(model.state_dict(), PATH)

И загрузите модель обратно на графический процессор/процессор:

model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device) # can be any device

Фактически, поскольку мы используем DDP для обучения модели, загрузка модели не зависит от количества машин, использованных для обучения предыдущей контрольной точки. Это позволяет нам горизонтально масштабировать парк Trn1 без изменений кода или негативных последствий для обучения модели. Эти контрольные точки на основе PyTorch можно использовать напрямую или даже с помощью скриптов torch для сценариев использования вывода на AWS Inferentia2 или других ускорителях.

Операционная стабильность

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

Как и в случае с существующим конвейером на базе графического процессора, мы добавили множество механизмов, обеспечивающих его эксплуатационную стабильность. Прежде чем начать обучение модели, мы проводим несколько тестов на работоспособность, чтобы оценить исправность машин. Эти тесты обычно включают простые тензорные операции для проверки работоспособности ускорительных устройств. Мы заметили, что для распределенного обучения важно также запускать тесты для проверки коллективного взаимодействия между экземплярами. Мы использовали набор тестов NCCOM из Neuron SDK для достижения этой цели, выполняя различные операции, такие как сбор всех, сокращение всех и уменьшение-разброс.

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

Благодаря этим механизмам устойчивости мы смогли достичь 98.5% успешности наших рабочих нагрузок на Trn1, что сопоставимо с нашими существующими показателями успешности конвейера графических процессоров.

Итоги

Чтобы проверить точность наших моделей, мы инициализировали две модели с одной и той же контрольной точки графического процессора и обучили одну на Trainium, а другую на сопоставимом графическом процессоре. Обе модели были обучены с использованием одних и тех же обучающих гиперпараметров. Набор данных, используемый для расчета метрик, представляет собой контрольный набор данных, и мы оцениваем точность модели на этом наборе данных каждые N глобальных шагов. Ось X — это глобальный шаг, а ось Y — точность модели. Мы наблюдали разницу в точности модели менее 0.1% в каждой точке на следующем графике.

Кроме того, чтобы оценить экономическую эффективность обучения модели, мы предпочитаем сравнивать время, необходимое для достижения сходимости модели. Мы считаем, что это дает более практичное представление об экономии затрат по сравнению с такими показателями, как стоимость токена, достигнутый FLOPS/доллар и другие факторы. Учитывая время обучения trn1.32xl и сопоставимых Эластичное вычислительное облако Amazon (Amazon EC2), мы заметили, что Trainium предлагает до 30% более низкую стоимость моделирования конвергенции.

Заключение

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

Благодаря нашему партнерству и сотрудничеству с командой Annapurna Neuron команда Amazon Search M5 смогла сэкономить до 30 % затрат, перейдя на Trainium. Команда может использовать Trainium и достигать качества моделей и равенства пропускной способности с сопоставимыми ускорителями на рынке. Совместимость контрольных точек и минимальные изменения кода с поддержкой XLA позволили M5 выбирать между несколькими ускорителями для своих рабочих нагрузок. Это позволило команде M5 воспользоваться преимуществами огромных вычислительных мощностей Trainium и создать независимые от ускорителей решения, чтобы порадовать клиентов Amazon.com. С операционной точки зрения Trainium доказал свою способность поддерживать сервисы уровня 1 в масштабе Amazon. Команда M5 продолжает переносить больше рабочих нагрузок в Trainium, чтобы предоставлять лучшие модели для Amazon с наименьшими затратами.

Подводя итог, можно сказать, что команде M5 удалось провести экономичное обучение машинному обучению промышленного уровня, добавив Trainium к парку ускорителей. Мы рекомендуем вам взглянуть на Trainium и другие устройства Neuron, такие как AWS Inferentia, чтобы воспользоваться преимуществами специально созданного кремния Amazon для рабочих нагрузок машинного обучения. Легко начните работу с помощью одного из многочисленных руководств, посвященных различным моделям, например Лама 2 доступна на Trainium..

Об авторах

Абхинандан Патни является старшим инженером-программистом в Amazon Search. Он занимается созданием систем и инструментов для масштабируемого распределенного обучения глубокому обучению и логических выводов в реальном времени.

Джеймс Парк — архитектор решений в Amazon Web Services. Он работает с Amazon.com над проектированием, созданием и развертыванием технологических решений на AWS, а также проявляет особый интерес к искусственному интеллекту и машинному обучению. В свободное время он любит искать новые культуры, новый опыт и оставаться в курсе последних технологических тенденций. Вы можете найти его на LinkedIn.

Джерри Маннил — инженер-программист в Amazon Search. Он работает над повышением эффективности, надежности и масштабируемости инфраструктуры распределенного обучения.

Кен Су — инженер-программист в Amazon Search. Он работает над повышением эффективности обучения и масштабируемым распределенным рабочим процессом обучения. Вне работы он любит пешие прогулки и теннис.

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