Обучение Qwen3 на AWS Trainium с помощью torchtitan
aws ·
В данной статье показано, как обучить модель Qwen3 на аппаратном обеспечении AWS Trainium с использованием фреймворка torchtitan. Авторы описывают три различных метода конфигурации: использование Hugging Face Transformers v5 в качестве бэкенда, применение нативной реализации Qwen3 от torchtitan и интеграцию ядер Neuron Kernel Interface (NKI) для моделей типа mixture-of-experts. Все подходы выполняются на одном экземпляре trn2.48xlarge с 64 логическими ядрами NeuronCores, используя стандартные интерфейсы PyTorch через TorchNeuron. В руководстве подробно изложены необходимые предварительные условия, включая требуемые версии программного обеспечения и настройку среды с помощью AWS Deep Learning AMI. Приведены примеры кода для дообучения на основе опубликованных весов, настройки механизмов внимания и подключения пользовательских ядер NKI через реестр переопределений. Особое внимание уделяется проверке стабильности обучения путем сравнения метрик потерь и обеспечению совместимости контрольных точек с форматами Hugging Face. Наконец, объясняются стратегии параллелизма, такие как шардирование по данным и тензорное параллелизм, для оптимизации использования ресурсов акселераторов.
aws