TorchNeuron: PyTorch върху чиповете на Amazon без CUDA

21.10

време за четене: 4 мин.

Amazon Web Services пусна TorchNeuron - инструмент, с който код на PyTorch върви върху собствените чипове Trainium почти без промени. От 20 октомври той е в публичен бета тест. Amazon атакува не чиповете на NVIDIA, а софтуерното заключване около тях.

В светла инженерна лаборатория робот пренася цял работещ модел от една сървърна стойка към съседна, като сменя само един кабел, а мъж програмист го гледа изненадано с дебел наръчник в ръце, докато жена колега вече проверява екрана на новата стойка.

Какво пусна Amazon

TorchNeuron е вграждане за PyTorch, широко използвана библиотека за обучение и изпълнение на AI модели. То добавя Trainium като пълноценно устройство, редом с видеокартите, и поддържа както незабавно изпълнение на операциите, така и компилиране на модела. Поддържа и разпределената работа на много чипове наведнъж, нужна при обучение на големи модели.

Според документацията на AWS инструментът е в публичен бета тест с версия 2.33 на софтуера Neuron и работи с второто и третото поколение Trainium. Hugging Face, която описва поддръжката в своите библиотеки, уточнява, че публичният бета тест започва на 20 октомври, и препоръчва вече да се работи директно с библиотеките Transformers, Diffusers и TRL вместо с досегашния междинен слой.

Amazon описва прехода като промяна на един ред в конфигурацията, без пренаписване на кода. Компанията посочва, че около 85% от AI приложенията са писани на PyTorch, но това е нейна оценка, не независимо измерване.

На практика промяната изглежда така: в Transformers моделът се насочва към новото устройство с параметъра device_map със стойност neuron, а при компилиране се избира съответният изпълнител. Обучаващите скриптове на TRL вървят върху Trainium с няколко специфични настройки. Изчисленията с осембитови числа с плаваща запетая, които пестят памет и време, изискват третото поколение на чипа.

Защо TorchNeuron е ход срещу CUDA

Предимството на NVIDIA не е само в чиповете. CUDA, софтуерният слой, през който приложенията стигат до видеокартите ѝ, е настроен години наред да работи гладко с PyTorch. Затова на практика кодът се пише за NVIDIA, а всеки друг хардуер изисква допълнителна работа по преноса.

TorchNeuron атакува точно тази работа. Досега пътят към Trainium минаваше през междинния слой PyTorch XLA и инструмента Optimum Neuron, а новото вграждане ги заменя с директна поддръжка в самата библиотека. Hugging Face отчита, че 89% от стоте най-използвани архитектури за генериране на текст вървят веднага, без допълнителна настройка, срещу 46% с предишния инструмент Optimum Neuron.

Amazon казва още, че ще внася промени в самата библиотека PyTorch, за да премахне кода, който я обвързва с конкретни чипове. Колин Брейс, вицепрезидент в Annapurna Labs, звеното на AWS за чипове, описва целта пред TechCrunch като премахване на заключването, което не оставя на разработчиците реален избор. PyTorch се управлява от Linux Foundation от 2022 г. и вече позволява пренос към чипове на AMD и Apple с подобна промяна в настройките.

Какво показва първият клиент

Най-конкретният пример е Deepgram, компания за разпознаване на реч. В технически текст на AWS моделът ѝ Nova-3 върху една инстанция trn2.3xlarge е обработил между 11% и 70% повече аудио на долар спрямо предишните ѝ изчислителни варианти при пакетна обработка. Сравнението е по цените към момента на изпитанието.

Тук числото зависи от източника. В TechCrunch Брейс цитира горната граница - 70% повече аудио на долар. Документът на AWS дава целия диапазон, а за решение на собствена фирма значение има долната граница, не рекордът. Точността е проверявана по процента грешно разпознати думи, но конкретни стойности не са публикувани.

Къде са границите

Документацията на AWS сама изброява ограниченията. Операция, за която няма реализация за Trainium, се изпълнява на процесора с предупреждение, което забавя работата. При компилиране такава операция спира процеса по подразбиране. Изчисленията с двойна точност вървят с единична, а всяка нова форма на данните се компилира наново при първото си ползване.

Това означава, че „един ред" важи за моделите, които вече се поддържат. За собствени или по-необичайни архитектури остава проверка операция по операция. Бета тестът е подходящ за пробно прехвърляне на едно натоварване, не за пълен преход на производствена система.

Наблюдаваме, че за фирмите, които наемат изчислителна мощ в облака, изборът на хардуер постепенно става въпрос на цена, а не на преписване на код. Колкото по-лесно се сменя чипът, толкова по-силна е позицията на клиента при договарянето.

Преди такова прехвърляне е разумно да се сравни не цената на час, а цената на свършена работа. Така мери и примерът с Deepgram - в аудио на долар, а не в часове работа на чипа.

Заключение

TorchNeuron не променя самите чипове, а сваля цената на преминаването между тях. Ако обещанието за един ред се потвърди извън бета теста, най-голямата промяна ще е в цените, които облачните клиенти могат да договорят.

Искате да започнете с AI от нулата? Безплатният курс „Основи на AI“ е 17 видео урока, около 35 минути, със сертификат: Започнете безплатния курс по AI

Помагаме на български екипи да изчислят реалната цена на AI натоварванията си и да сравнят облачните варианти. Услугите ни покриват одит, оценка на готовност, обучение на екипи и имплементация на специализирани AI автоматизации за всяка организация. Препоръчваме консултация, преди да се прехвърли производствен модел към нов хардуер. [Свържете се с нас за консултация: academy@razvivai.se]