Крупнейшие компании, занимающиеся разработкой в области машинного обучения, представили проект 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 (Intermediate Representation Execution Environment) - компилятор и runtime, преобразующий модели машинного обучения в универсальное промежуточное представление, основанное на формате MLIR (Multi-Level Intermediate Representation) от проекта LLVM. Из особенностей отмечается возможность предварительной компиляции (ahead-of-time), поддержка управления потоком, возможность использования динамических элементов в моделях, оптимизация для разных CPU и GPU, низкие накладные расходы.

Основные преимущества инструментария OpenXLA:

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