Entraîner Qwen3 sur AWS Trainium avec torchtitan
aws ·
Cet article montre comment entraîner le modèle Qwen3 sur du matériel AWS Trainium en utilisant le framework torchtitan. Les auteurs présentent trois méthodes de configuration distinctes : l'utilisation de Hugging Face Transformers v5 comme backend, l'exploitation de l'implémentation native de Qwen3 dans torchtitan, et l'intégration de noyaux Neuron Kernel Interface (NKI) pour les modèles à experts multiples. Toutes ces approches s'exécutent sur une seule instance trn2.48xlarge disposant de 64 cœurs Neuron logiques, en tirant parti des interfaces PyTorch standard via TorchNeuron. Le guide détaille les prérequis spécifiques, notamment les versions logicielles requises et la configuration de l'environnement à l'aide de l'AMI Deep Learning d'AWS. Il fournit des exemples de code pour le fine-tuning à partir de poids publiés, la personnalisation des mécanismes d'attention et l'attachement de noyaux NKI définis par l'utilisateur via un registre de remplacement. La démarche met l'accent sur la vérification de la stabilité de l'entraînement en comparant les métriques de perte et en assurant la compatibilité des points de contrôle avec les formats Hugging Face. Enfin, elle explique les stratégies de parallélisme, telles que le sharding par données et le parallélisme tensoriel, afin d'optimiser l'utilisation des ressources entre les accélérateurs.
aws