На одном trn2.48xlarge обучение Qwen3 в torchtitan собирается тремя способами: через Hugging Face Transformers v5 как modeling backend, через нативную реализацию Qwen3 и через NKI-ядра mixture-of-experts из рецептов torchtitan-neuron. Запуск у всех трёх один и тот же: torchrun и общая структура конфигов. Меняется backend модели и набор ядер.
Причина, по которой это работает без отдельного форка под AWS, в TorchNeuron. Он регистрирует Trainium как нативное устройство PyTorch, поэтому torchtitan видит ускоритель при старте и общается с ним через стандартные интерфейсы PyTorch.
Практический смысл: если у вас уже есть пайплайн torchtitan или код на Transformers, перенос на Trainium не требует переписывать цикл обучения. Достаточно поменять конфиг, закрепить версии стека Neuron и убедиться, что веса загрузились корректно.
Что даёт связка torchtitan + TorchNeuron на Trainium
После установки torch_neuronx устройство Trainium появляется в PyTorch под именем neuron, и torchtitan обнаруживает его при запуске. Отдельного форка под AWS нет: torchtitan обращается к Trainium через стандартные интерфейсы PyTorch, а компиляцию и раскладку по NeuronCores берёт на себя стек Neuron. Walkthrough AWS по обучению Qwen3 на Trainium строит весь сценарий именно на этом механизме.
Режим компиляции выбирают под задачу. В eager-режиме Neuron компилирует операции по мере выполнения: удобно отлаживать загрузку весов и проверять формы тензоров. С torch.compile компилируется захваченный граф целиком, путь исполнения получается быстрее, но ошибки всплывают позже и читать их сложнее.
Конфигурации покрывают оба сценария обучения: предобучение со случайной инициализации и дообучение опубликованных весов. Их же берут как заготовку, когда нужно заменить реализацию модели или отдельные ядра под свою нагрузку.
Три способа на одном железе: чем именно они отличаются
| Способ | Как подключается | Что даёт | Когда брать |
|---|---|---|---|
| 1. Transformers v5 | AutoModelForCausalLM как modeling backend | Прямое чтение и запись опубликованных safetensors, минимум своего кода | Дообучение готовых весов, знакомая кодовая база |
| 2. Нативная Qwen3 | Изменение config.model_spec | Контроль над структурой модели внутри torchtitan | Планируется кастомизация модели и ядер |
| 3. NKI MoE | Ядра из рецептов torchtitan-neuron плюс override registry | Ускоренные MoE-ядра и подмена отдельной операции своим ядром | Эксперименты с NKI и mixture-of-experts |
Общая часть у всех трёх одна: конфигурация это Python-функция, которая стартует от базового конфига и обновляется на каждом шаге. Поэтому переключение между путями выглядит как правка конфига, а не как новый пайплайн.
Что нужно подготовить: инстанс, окружение и версии
Walkthrough рассчитан на один trn2.48xlarge. Внутри 16 ускорителей Trainium2 и 128 физических NeuronCores, при LNC=2 они делятся на 64 логических NeuronCores. Это число задаёт верхнюю границу для раскладки параллелизма.
Окружение берут одно из двух: AWS Deep Learning AMI for Neuron либо контейнер Neuron PyTorch training. Оба варианта уже содержат драйвер, runtime и компилятор Neuron, без которых обучение на Trainium не запустится, так что собирать стек с нуля не нужно.
| Компонент | Версия в walkthrough |
|---|---|
| torch | 2.14.0 |
| torch_neuronx | 2.14.3 |
| transformers | 5.16.0 |
| torchtitan-neuron | release-2.33 (0.2.2+neuron.1) |
Перед установкой torchtitan-neuron проверяют, нет ли его в образе:
python3 -c "import torchtitan, torchtitan_neuron; print(torchtitan.__version__, torchtitan_neuron.__file__)"Команда печатает версию и путь к пакету. Если torchtitan_neuron уже импортируется, шаг установки пропускают. Если нет, клонируют репозиторий torchtitan-neuron, ставят зависимости из requirements.txt, а сам пакет и остальные зависимости torchtitan устанавливают с флагом --no-deps. Флаг здесь не формальность: без него pip может подтянуть обычную сборку torch и затереть Neuron-вариант, после чего устройство neuron исчезнет.
Сведения о требованиях к окружению в первоисточнике ограничены одним trn2.48xlarge и этими двумя вариантами AMI или контейнера. Для другой конфигурации железа шаги придётся адаптировать. Полный список команд приведён в исходном разборе AWS.
Загрузка весов Qwen3 с Hugging Face Hub
Артефакты модели забирают с Hugging Face Hub в локальный каталог, а путь к нему передают в конфигурацию обучения:
hf download Qwen/Qwen3-4B-Instruct-2507 --local-dir "$PWD/assets/hf/Qwen3-4B-Instruct-2507"В первой конфигурации эти safetensors читаются и записываются напрямую, без промежуточного формата конвертации. Контекст партнёрства Hugging Face и AWS вокруг обучения на Trainium и Inferentia разобран отдельно: Hugging Face и AWS: что даёт партнёрство.
Способ 1: Transformers v5 как modeling backend
Здесь torchtitan использует Transformers v5 как backend модели: Qwen3 поднимается через AutoModelForCausalLM, а опубликованные safetensors читаются и пишутся напрямую. Для команды с готовым кодом на Transformers это самый короткий путь: архитектуру не переписывают, класс модели берут из библиотеки.
Конфиг остаётся Python-функцией, которая начинается от базового конфига и обновляется по шагам. В него передают путь к локальному каталогу с весами, версию Transformers закрепляют: в walkthrough это 5.16.0.
Ограничение прямолинейное: контроль над ядрами минимален. Вы получаете то, что стек Neuron делает для операций Transformers, и не можете заменить отдельную операцию своим ядром, не переходя на другой путь. Для дообучения опубликованных весов этого достаточно, для экспериментов с NKI уже нет.
Способ 2: нативная реализация Qwen3 в torchtitan
Переключение делается правкой config.model_spec: вместо внешнего backend torchtitan берёт собственную реализацию Qwen3. Launch-паттерн не меняется, конфиг остаётся той же Python-функцией.
Структура модели описана внутри torchtitan, поэтому её проще менять: головы, нормализации, порядок операций. К этой же базе подключаются свои ядра. Такой вариант занимает середину между быстрым стартом на Transformers и работой с NKI напрямую.
Цена решения в том, что расхождение с эталонной реализацией из Transformers контролируете вы. Если дообучаете опубликованные веса, проверьте совпадение имён и раскладки параметров до запуска: рассинхрон проявится сразу на первом шаге лосса.
Способ 3: NKI-ядра MoE и своё ядро через override registry
Третий путь начинается с рецептов torchtitan-neuron: оттуда берут NKI-ядра mixture-of-experts. Дальше override registry torchtitan позволяет подключить собственное ядро, то есть подменить реализацию конкретной операции, не трогая остальной граф.
Пример в walkthrough: ядро rope_hf из nkilib с регистрацией autograd. Регистрация нужна потому, что NKI-ядро считает только forward, а обучению требуется корректный backward. Вокруг ядра описывают autograd-функцию, которая возвращает градиенты в понятном PyTorch виде. Пропустить этот шаг нельзя: без него обратный проход либо упадёт, либо посчитает градиенты неверно.
Ограничение жёсткое: кастомное ядро требует длины последовательности, кратной 128 * LNC. При LNC=2 это 256 токенов, и последовательность длиной 200 или 300 токенов такое ядро не обслужит. Длину батча, packing и сплиты стоит планировать заранее. Детали рецепта и конфига собраны в материале AWS с тремя конфигурациями.
Ограничения MoE-примера на одной инстансе
MoE-пример урезан до 4 слоёв и 64 экспертов, чтобы уместиться на одном trn2.48xlarge. Это демонстрация механики: как подключаются ядра, как считается forward и backward, как выглядит конфиг. Продакшн-конфигурацией такой пример не считается. Переносить его поведение на полноразмерный MoE с десятками слоёв и сотнями экспертов не стоит: там понадобится больше памяти и другой план параллелизма.
Запуск через torchrun и проверка первого шага лосса
Все три конфигурации стартуют одинаково: torchrun, стандартные аргументы torchtitan и путь к конфигу. Схематично запуск выглядит так, конкретный набор флагов зависит от версии torchtitan:
torchrun <аргументы torchtitan> --job.config_file <путь_к_конфигу>Меняются путь к конфигу и его содержимое: modeling backend, путь к весам, степени параллелизма. Скрипт запуска и раскладка процессов остаются прежними, поэтому переключение между способами занимает минуты.
Первый шаг лосса служит самой дешёвой проверкой того, что опубликованные веса действительно загрузились. При корректных safetensors и конфиге значение на первом шаге близко к ожидаемому для выбранной модели и дальше плавно снижается. Аномально высокий или скачущий лосс чаще всего указывает на одно из трёх: не тот каталог с весами, рассинхрон имён параметров, неподходящий конфиг модели. Точные числа зависят от модели, размера батча и режима компиляции, поэтому ориентир в walkthrough задан по поведению, а не по эталонной цифре.
Параллелизм: data_parallel_shard_degree, tensor_parallel_degree и FSDP на DeviceMesh
torchtitan раскладывает обучение по многомерной сетке устройств, DeviceMesh. Разные оси отвечают за разные виды параллелизма, а FSDP шардирует параметры, состояние оптимизатора и градиенты внутри data-parallel оси. Степени задают два параметра конфига.
- data_parallel_shard_degree определяет, на сколько шардов FSDP делит модель и состояние оптимизатора. Больше шардов означает меньше памяти на устройство и больше обменов между ними.
- tensor_parallel_degree определяет, на сколько частей режутся сами тензоры слоёв. Это спасает, когда один слой не помещается в память NeuronCore, но добавляет обмены на каждом слое.
Оба значения привязаны к железу. На trn2.48xlarge доступно 16 ускорителей Trainium2 и до 64 логических NeuronCores при LNC=2, поэтому произведение степеней должно укладываться в доступное число логических ядер без остатка. Рабочий порядок такой: сначала подбирают data_parallel_shard_degree, чтобы модель и оптимизатор помещались в память, и только потом добавляют tensor_parallel_degree, если памяти всё равно не хватает.
От LNC зависит и число логических ядер, и требование к длине последовательности в кастомных NKI-ядрах (кратность 128 * LNC). Смена LNC меняет и то, и другое, поэтому батчи придётся пересчитать. Как устроен FSDP и распределённый запуск в соседних библиотеках, разобрано в материале про Accelerate 1.0 и распределённое обучение, хотя torchtitan использует собственный слой поверх PyTorch.
Кому это нужно и где границы применимости
Сценарий оправдан, когда инфраструктура уже живёт в AWS и стоит задача обучить или дообучить Qwen3 на Trainium с контролем над реализацией модели и ядер. Три конфигурации дают постепенное погружение: быстрый старт на Transformers v5, затем нативная модель, затем NKI-ядра MoE.
Границы стоит держать в голове до запуска:
- одна инстанса trn2.48xlarge, масштабирование за её пределы в walkthrough не описано;
- MoE-пример урезан до 4 слоёв и 64 экспертов, это демо механики;
- кастомное NKI-ядро требует длины последовательности, кратной 128 * LNC;
- стек привязан к версиям: torch 2.14.0, torch_neuronx 2.14.3, transformers 5.16.0, torchtitan-neuron release-2.33;
- интеграция Transformers v5 и Trainium-специфичные override-модули меняются между релизами torchtitan, поэтому версии закрепляют.
Если задача сводится к инференсу или коротким экспериментам, вход в Trainium потребует Neuron-сборок, флага --no-deps и возни с ограничениями ядер без явной отдачи. Для контролируемого training-пайплайна внутри AWS три конфигурации дают рабочую отправную точку.
Начните с одной инстансы, закрепите перечисленные версии, прогоните способ 1 на Qwen3-4B-Instruct-2507 и посмотрите на первый шаг лосса. Дальше переключайте config.model_spec и подключайте NKI-ядра по одному, сверяя лосс после каждого изменения.