Представен е 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. Из особенностей стоит отметить возможность предварительной компиляции (ahead-of-time), поддержку управления потоком, возможность использования динамических элементов в моделях, оптимизацию для разных 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