PyTorch a annoncé TorchNeuron, un backend natif qui permet à PyTorch de reconnaître AWS Trainium comme un périphérique standard. Cette nouvelle bibliothèque permet aux développeurs de transférer un modèle vers Trainium avec la simple commande .to('neuron'), puis d’utiliser torch.compile pour obtenir des performances adaptées à la production. Elle prend en charge l’exécution eager, le mode eager adaptatif et les noyaux personnalisés via le Neuron Kernel Interface. Le code sera publié en open source sous licence Apache 2.0, et la bêta publique débutera le 20 oct.
Pourquoi c'est important
Les chercheurs peuvent exploiter les performances de Trainium avec le même code PyTorch, sans avoir à réécrire leurs flux de travail.