AWS Neuron представляет поддержку Neuron Kernel Interface (NKI), NxD Training и обучения на платформе JAX

Проведено: 17 сент. 2024 г.

Сегодня AWS объявляет о выпуске пакета средств разработки ПО Neuron 2.20, в котором впервые представлена бета-версия интерфейса программирования Neuron Kernel Interface (NKI) для AWS Trainium и Inferentia. Этот интерфейс позволяет разработчикам создавать оптимизированные вычислительные ядра для новых функциональных возможностей, оптимизации и научных инноваций. Кроме того, в этом выпуске представлена бета-версия библиотеки NxD Training на основе PyTorch для эффективного распределенного обучения с удобным совместимым с NeMo интерфейсом. В этом выпуске также впервые реализована поддержка платформы JAX (бета-версия).

AWS Neuron — это пакет средств разработки ПО (SDK) для инстансов на базе AWS Inferentia и Trainium, специально созданный для генеративного искусственного интеллекта. Neuron интегрируется с популярными платформами машинного обучения, такими как PyTorch. Он включает компилятор, среду выполнения, инструменты и библиотеки для поддержки высокопроизводительного обучения и логического вывода в моделях ИИ на инстансах Trn1 и Inf2.

В этом выпуске добавлены функции поддержки и улучшена производительность обучения моделей и логического вывода. Для обучения в этом выпуске добавлена поддержка моделей 8B и 70B из Llama 3.1 с длинами последовательностей до 32K, а также поддержка torch.autocast() для смешанной точности в PyTorch и поддержка методов PEFT LoRA. Для логического вывода в Neuron 2.20 добавлена поддержка моделей Llama 3.1 (405b, 70b, 8b) и моделей диффузионных трансформеров (DiT), таких как Pixart-alpha и Pixart-sigma. Кроме того, в этом выпуске добавлена поддержка логического вывода с выборкой top-p на устройстве, контекста длиной 128K и метода Flash Decoding. Также в этом выпуске появилась поддержка операционной системы Rocky 9.0 и операторов RMSNorm и RMSNormDx в компиляторе Neuron.

Дополнительные сведения см. в примечаниях к выпуску Neuron.