Як Amazon Search M5 заощадив 30% на вартості навчання LLM за допомогою AWS Trainium

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

читають: 0

Десятиліттями Amazon впроваджував інновації в машинному навчанні (ML), надаючи чудові враження своїм клієнтам. З перших днів Amazon використовував машинне навчання для різних випадків використання, таких як рекомендації книг, пошук і виявлення шахрайства. Подібно до решти індустрії, досягнення прискореного апаратного забезпечення дозволили командам Amazon розробляти модельні архітектури з використанням нейронних мереж і глибокого навчання (DL).

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

Як і багато організацій ML, прискорювачі в основному використовуються для прискорення навчання DL та висновків. Коли AWS запустила спеціально створені прискорювачі з першим випуском AWS Inferentia у 2020 році команда M5 швидко почала використовувати їх для більш ефективного розгортання виробничих робочих навантажень, заощаджуючи витрати та скорочуючи затримку. Минулого року AWS запустила його AWS Trainium прискорювачі, які оптимізують продуктивність за ціною для розробки та створення моделей DL наступного покоління. У цій публікації ми обговорюємо, як M5 вдалося зменшити витрати на навчання своїх моделей на 30%, і ділимося деякими з найкращих практик, які ми навчилися на цьому шляху.

Примірники Trainium

Завдяки прогресу в спеціально створених прискорювачах Amazon також пропонує переконливі прискорювачі у вигляді AWS Inferentia та Trainium. Як випливає з назви, ці мікросхеми оптимізовані, щоб перевершити потреби висновків і навчальних навантажень відповідно. Для широкомасштабного навчання моделей фундаментів, розмір яких досягає мільярдів параметрів, Trainium Екземпляри Trn1 і Trn1n є ідеальним вибором завдяки своїм характеристикам. Екземпляри Trn1 оснащені найсучаснішим обладнанням NeuronCore-v2, і мають велику кількість прискорювача обчислень і пам’яті. Екземпляри Trn1n також можна вибрати для більшої пропускної здатності мережі (1,600 Гбіт/с), тому вони ідеально підходять для продуктивного навчання з оптимізацією витрат.

Щоб використовувати прискорювачі, вам потрібен програмний рівень для їх підтримки. З чіпами Trn і Inf, AWS Neuron SDK розблоковує спеціальні прискорювачі Amazon за допомогою PyTorch XLA. PyTorch XLA перетворює нетерплячий режим PyTorch на ліниву реалізацію на основі графіків. Потім ці графіки використовуються та компілюються для використання з прискорювачем. PyTorch Neuron (частина Neuron SDK) дозволяє користувачам PyTorch навчати свої моделі на Trainium NeuronCores за допомогою кількох рядків коду.

Модель і навантаження

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

Інфраструктура для використання виробничого трубопроводу Пакет AWS з стратегії справедливого розподілу в черзі, використовуючи багатовузловий кластер trn1.32xlarge із підтримкою EFA як обчислення для навчання моделі. Функціонально виробничий конвеєр виконує поступове навчання моделі, оцінку навченої моделі та офлайн-пакетний висновок на навченій моделі, використовуючи PyTorch як базову бібліотеку DL.

Цілі

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

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

У наступних розділах ми поділимося нашим шляхом роботи у зворотному напрямку від цього критерію та нашими знаннями щодо підтримки робочих навантажень виробництва в масштабі 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

Традиційно цикл розробки моделі з графічним процесором передбачає внесення змін до моделі або сценарію навчання та безпосередній запуск на пристрої графічного процесора. Для таких прискорювачів, як Trainium, які використовують XLA, потрібен додатковий етап, перш ніж навчання моделі можна буде запустити на прискорювачі. Графіки обчислень XLA можна запускати лише після їх компіляції. Як правило, існує два способи виконати цю компіляцію: завчасно (AOT), коли ви спочатку відстежуєте та компілюєте всі графіки, а потім запускаєте їх, або Just In Time (JIT), коли графіки трасуються, компілюються та запускаються, коли вони зустрічаються. Neuron SDK надає обидва з коробки. Як правило, спочатку виконується компіляція AOT. Потім після цієї компіляції запускаються графіки. Якщо зустрічаються нові графіки, середовище виконання Neuron викликає компіляцію JIT перед їх запуском. Для компіляції AOT передбачено Neuron SDK neuron_parallel_compile, утиліта компіляції, яка витягує графіки з пробного запуску сценарію навчання та виконує паралельну компіляцію AOT.

Важливим аспектом компіляції AOT є забезпечення того, щоб під час навчання не створювалися нові графіки обчислень. Одним із джерел нових обчислювальних графів (і, отже, повторних компіляцій) є динамічні форми навчальних пакетів під час навчання моделі. Ми виявили, що використання статичних форм і пакетів фіксованого розміру усуває компіляції часу навчання та значно покращує пропускну здатність навчання без жодного впливу на точність моделі. Накладаючи такі обмеження на навчання, ми помітили, що лише 4–5 кроків навчання моделі, один крок перевірки моделі та одноразове встановлення контрольних точок моделі необхідні для трасування всіх графіків під час компіляції AOT. Важливо зазначити, що 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. Це було критично важливо для того, щоб ми могли ініціалізувати нашу модель за допомогою найкращої попередньо навченої моделі без будь-яких простоїв виробництва чи втрати точності моделі.

Оскільки модель GPU було збережено за допомогою стандартних утиліт збереження моделі PyTorch, ми змогли використати утиліту завантаження контрольних точок PyTorch, щоб завантажити модель GPU на пристрої 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)

І завантажте модель назад на GPU/CPU:

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

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

Стабільність роботи

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

Як і в існуючому конвеєрі на основі GPU, ми додали численні механізми, щоб зробити конвеєр стабільним у роботі. Перед початком навчання моделі ми проводимо кілька тестів на працездатність, щоб оцінити працездатність машин. Ці тести зазвичай включають прості тензорні операції для перевірки справності пристроїв прискорювача. Ми помітили, що для розподіленого навчання важливо також запускати тести для перевірки спільного зв’язку між примірниками. Ми використовували Набір тестів NCCOM з Neuron SDK, щоб досягти цього, виконуючи різноманітні операції, такі як all-gather, all-reduce та reduce-scatter.

Навіть після виконання згаданих пропозицій ми помітили, що тимчасові проблеми неминучі в будь-якому конвеєрі, незалежно від основного прискорювача. Щоб підвищити стійкість будь-якого навчального конвеєра, ми рекомендуємо створити механізми повторних спроб для вирішення цих потенційних проблем. Ми використовуємо Пакетні автоматичні повтори AWS щоб повторити завдання, які стикаються з тимчасовою помилкою під час навчання моделі. Ці перезапуски можуть бути дорогими, якщо збій виникне наприкінці навчання. Щоб протистояти цій проблемі, ми адаптували наші навчальні сценарії для завантаження попередньо навченої контрольної точки моделі та продовження навчання з цієї точки. Завдяки цій функції ми можемо агресивно перезапускати невдалі навчальні завдання з мінімальними витратами.

За допомогою цих механізмів стійкості ми змогли досягти 98.5% успіху для наших робочих навантажень на Trn1, що можна порівняти з нашим існуючим показником успішності конвеєра GPU.

результати

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

Крім того, щоб оцінити економічну ефективність навчання моделі, ми віддаємо перевагу порівнянню часу настінного годинника, необхідного для досягнення конвергенції моделі. Ми вважаємо, що це дає більш практичний погляд на економію коштів порівняно з такими показниками, як вартість токена, досягнутий FLOPS/долар та інші фактори. Враховуючи час навчання trn1.32xl і порівнянний Обчислювальна хмара Amazon Elastic (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 для робочих навантажень ML. Почніть легко з одного з багатьох навчальних посібників із різними моделями, наприклад Llama 2, доступний на Trainium.

Про авторів

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

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

Джеррі Менніл є інженером програмного забезпечення в Amazon Search. Він працює над підвищенням ефективності, надійності та масштабованості інфраструктури розподіленого навчання.

Кен Су є інженером програмного забезпечення в Amazon Search. Він працює над підвищенням ефективності навчання та масштабованим розподіленим процесом навчання. Поза роботою він любить походи та теніс.

RJ є інженером в Amazon. Він будує та оптимізує системи для розподілених систем для навчання та працює над оптимізацією адаптаційних систем для зменшення затримки для ML Inference. Поза роботою він досліджує використання Generative AI для створення рецептів їжі.