Представен е OpenXLA, инструментариум за оптимизация и компилация на модели за машинно обучение

Най-големите компании, които се занимават с разработки в сферата на машинното обучение, представиха проекта OpenXLA, насочен към съвместното развитие на инструменти за компилация и оптимизация на модели за системи за машинно обучение. Под крилото на проекта премина разработката на инструменти, позволяващи унификация на компилацията на модели, създадени във фреймворковете TensorFlow, PyTorch и JAX, за ефективно обучение и изпълнение на различни GPU и специализирани ускорители. В съвместната работа по проекта се включиха компании като Google, NVIDIA, AMD, Intel, Meta, Apple, Arm, Alibaba и Amazon.

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

OpenXLA се състои от три основни компонента, чийто код се разпространява под лицензия Apache 2.0:

• XLA (Accelerated Linear Algebra) — компилатор, позволяващ оптимизация на модели за машинно обучение за високопроизводително изпълнение на различни хардуерни платформи, включително GPU, CPU и специализирани ускорители от различни производители.
• StableHLO — спецификация и базова реализация на набор от високопрофилни операции (HLO, High-Level Operations) за използване в модели на системи за машинно обучение. Той служи като междинен слой между фреймворковете за машинно обучение и компилаторите, преобразуващи модела за изпълнение на конкретно оборудване. Междинни слоеве за генериране на модели във формата StableHLO са подготвени за фреймворковете PyTorch, TensorFlow и JAX. Базата за StableHLO е наборът MHLO, който е разширен с поддръжка за серийно предаване и версиониране.
• IREE (Среда за изпълнение на междинно представяне) — компилатор и runtime, който преобразува модели на машинно обучение в универсално междинно представяне, базирано на формата MLIR (Многостепенно междинно представяне) от проекта LLVM. Сред особеностите се отличава възможността за предварителна компилация, поддръжка на управление на потока, възможност за използване на динамични елементи в моделите, оптимизация за различни CPU и GPU, ниски разходи.

Основни предимства на инструмента OpenXLA:

• Постигане на оптимална производителност без необходимост от дълбочинно потапяне в кодиране, специфично за определени устройства. Предоставяне на готови оптимизации, които включват опростяване на алгебрични изрази, ефективно разположение в паметта, планиране на изпълнението с оглед на намаляване на пиковото потребление на памет и натоварвания.
• Оптимизация на мащабирането и разпаралеляването на изчисленията. На разработчика е достатъчно да добави анотации за подмножество критични тензори, на базата на които компилаторът може автоматично да генерира код за паралелни изчисления.
• Осигуряване на преносимост чрез поддръжка на различни хардуерни платформи, като GPU AMD и NVIDIA, CPU на базата на архитектури x86 и ARM, ML-ускорители TPU на Google, IPU AWS Trainium, Inferentia, Graphcore и Cerebras Wafer-Scale Engine.
• Поддръжка на свързване на разширения с внедряване на допълнителни функции, като поддръжка за писане на примитиви за дълбоко машинно обучение с използване на CUDA, HIP, SYCL, Triton и други езици за паралелни изчисления. Възможност за ръчно настройване на тесни места в моделите.

Източник: opennet.ru