Подробный разбор · Около 6 мин чтения
Qwen3 на Trainium: три пути torchtitan для проверки весов и ядер
Руководство AWS от 21 октября объединяет Transformers, нативный Qwen3 и ядра NKI на одном экземпляре Trainium2, но не доказывает преимущества в скорости или стоимости.
- Опубликовано здесь
- Дата источника
- Источник проверен
Техническое руководство AWS от 21 октября описывает обучение Qwen3 через torchtitan на Trainium в трёх конфигурациях: с бэкендом Transformers v5, нативной моделью torchtitan и уменьшенным примером смеси экспертов с ядрами NKI. Полезное новшество — конкретная схема интеграции, включая проверки контрольных точек и численных результатов. Это не выпуск новой модели Qwen и не бенчмарк, доказывающий удешевление обучения.
Слушать новости
Прослушайте эту статью или включите новости подряд.
Подготовка чтения…
Голос устройства
Речь обрабатывает браузер или устройство. AISubHub не отправляет текст в аудио-API. Доступные голоса и работа при заблокированном экране зависят от устройства.
При продолжении короткая фраза может повториться.
Кратко
- В руководстве используется один trn2.48xlarge: 16 ускорителей Trainium2, 128 физических NeuronCore и 64 логических при LNC=2. Команда запуска создаёт 64 процесса обучения.
- Плотные примеры используют файлы Qwen3-4B-Instruct-2507 и десять шагов на c4_test. Пример MoE сокращён до четырёх слоёв и 64 экспертов, с выбором четырёх экспертов на токен; полная неизменённая модель 30B-A3B здесь не обучается.
- AWS сообщает о функции потерь на первом шаге 6.54 при загрузке опубликованных весов и 12.43 при случайной инициализации. Проверка собственного вращательного позиционного кодирования с входами float32 дала максимальную абсолютную разницу 0.0. Эти числа проверяют разные свойства.
Одна интеграция устройства и три места для замены реализации модели
TorchNeuron регистрирует Trainium в PyTorch как нативное устройство neuron, позволяя torchtitan использовать привычные интерфейсы. Первый путь поручает построение модели AutoModelForCausalLM из Transformers v5. Второй выбирает собственную реализацию Qwen3 в torchtitan через config.model_spec. Третий подключает специфичные для Trainium операторы NKI к слоям смеси экспертов и показывает пользовательскую замену. Это варианты реализации за общей точкой входа в обучение, а не три новых модели Qwen.
Руководство фиксирует torch 2.14.0, torch_neuronx 2.14.3, Transformers 5.16.0 и torchtitan-neuron release-2.33 версии 0.2.2+neuron.1. Нужны Neuron Deep Learning AMI или обучающий контейнер с драйвером, средой выполнения и компилятором. По нашей оценке, привычный код PyTorch сужает поверхность интеграции, но не делает проверенное сочетание зависимостей необязательным. AWS прямо предупреждает, что интеграция Transformers и модули замен могут меняться между выпусками torchtitan.
Рецепт 4B проверяет инициализацию и переносимость контрольной точки
Первая конфигурация указывает Qwen/Qwen3-4B-Instruct-2507, использует c4_test, задаёт десять шагов и локальный размер пакета один, а состояние параллелизма данных распределяет между 64 запущенными рангами. Загрузка архитектуры и загрузка опубликованных весов — разные операции. Конфигурация дообучения явно включает чтение контрольной точки из скачанных файлов и экспортирует итоговую модель в Hugging Face safetensors. Заявленная функция потерь первого шага составляет 6.54 с опубликованными весами против 12.43 при случайной инициализации.
Разница помогает проверить инициализацию, но не доказывает улучшение ассистента за десять шагов или окончательную сходимость. AWS также повторно загружает экспортированную контрольную точку через Transformers для проверки ключей, учитывая lm_head.weight, связанный с эмбеддингами. Нативный путь начинается с конструктора qwen3_32b, поскольку в проверенной сборке нет общего конструктора 4B; затем спецификация заменяется на 4B, а глобальный размер пакета устанавливается в 64. Название конструктора поэтому не означает обучение модели 32B.
Конфигурация MoE намеренно меньше, чем предполагает имя семейства
Рецепт Qwen3-30B-A3B сокращён до четырёх слоёв и 64 экспертов, чтобы пример поместился на одном экземпляре. top_k=4 направляет каждый токен к четырём экспертам. Источник использует шардирование параллелизма данных на 64 рангах, степень тензорного параллелизма один и не использует экспертный параллелизм. Число экспертов — параметр модели, а не число устройств экспертного параллелизма. Это различие не позволяет превратить небольшую демонстрацию интеграции в заявление об обучении исходной полной модели.
Пример работает в режиме eager с выключенной компиляцией. Зарегистрированные операторы NKI также поддерживают torch.compile, но смена режима выполнения — отдельное решение. В проверенной сборке рецепт заменяет блок MoE ядрами сгруппированных экспертов и выбирает внимание на основе масштабированного скалярного произведения и вращательное позиционное кодирование без gather. AWS объясняет, как подключить autograd к операторам через сведения, сохранённые при прямом проходе. Сравнения пропускной способности, которое доказывало бы ускорение этой конкретной конфигурации, не приводится.
Собственное ядро требует проверки расположения данных, градиентов и области замены
Пример вращательного позиционного кодирования показывает, что регистрация сложнее простого вызова ядра. torchtitan хранит query и key в порядке пакет, последовательность, головы, размерность головы; rope_hf ожидает пакет, головы, последовательность, размерность головы. Обёртка меняет расположение, выделяет память для результата и выполняет обратное преобразование. Отдельный оператор обратного прохода использует сохранённый вращательный кэш, а для самого кэша градиент не вычисляется. Длина последовательности должна быть кратна 128, умноженным на LNC: при LNC два — кратна 256.
Затем реестр замен подменяет конфигурацию нужного модуля. Стартовые журналы должны показывать одну замену nki_rope на каждый слой Qwen3; более узкие шаблоны имён позволяют выбрать отдельные слои. AWS сообщает максимальную абсолютную разницу 0.0 для query и key при сравнении float32 с реализацией PyTorch. Это подтверждает конкретный численный тест, но не все типы данных, формы или обучающие запуски. Статья прямо называет пример интеграцией ядра и autograd, а не улучшением производительности.
Как превратить руководство в измеряемый пилот
Первый шаг включает компиляцию и занимает больше времени, чем последующие. Журналы показывают функцию потерь, норму градиента, использование высокоскоростной памяти, пропускную способность и загрузку FLOPs модели, но руководство не даёт контролируемого сравнения полной стоимости. С ростом модели, длины последовательности или локального пакета память на ранг остаётся практическим ограничением. Тензорный параллелизм можно объединять с полностью шардированным параллелизмом данных на многомерном DeviceMesh; показанную настройку MoE нельзя выдавать за такой более крупный эксперимент.
AISubHub предлагает сохранить проверенные версии, проверить выбранную модель и 64 ранга, убедиться в загрузке и экспорте контрольной точки и сравнить выходы собственного ядра до длительного запуска. Затем следует измерить прогрев отдельно от устойчивой работы и оценить сходимость на целевых данных. Нужные контрольные точки необходимо сохранить до остановки или удаления арендованного оборудования. Это выведенные из руководства этапы оценки, а не воспроизведённые нами результаты. Они не устанавливают права по подписке, бесплатную квоту обучения или универсальное преимущество Trainium.
Почему это важно
Анализ AISubHub: команды могут отдельно проверить совместимость модели, переносимость контрольных точек и корректность ядер до запуска полноценного обучения. Общая структура запуска уменьшает различия между рецептами, но не отменяет требований к версиям, памяти, компиляции и численному поведению.
Ограничения и доступность
Это техническое руководство сотрудников AWS Josh Longenecker и Alex Hoffman, а не аппаратный тест AISubHub. В нём нет сопоставимой полной стоимости обучения, итогового результата сходимости или доказанного ускорения от собственного ядра. Доступность и оплата экземпляра, точные версии ПО и проверка рабочей нагрузки остаются ответственностью оператора. Преимущества пользовательских подписок не объявлены.
Источники и подтверждения
Дата источника: · Источник проверен:
Материалы подготовлены с помощью ИИ и сверены с указанными первоисточниками. Заявления поставщиков атрибутированы; публикация не означает рекомендацию.