OpenXLA — это совместно разработанная экосистема компилятора машинного обучения с открытым исходным кодом.
В последнее время, представлены крупнейшие компании, занимающиеся разработками в области машинного обучения проект OpenXLA, предназначенный для совместной разработки инструментов компилировать и оптимизировать модели для систем машинного обучения.
Проект взял на себя разработку инструментов, позволяющих унифицировать компиляцию моделей, подготовленных в фреймворках TensorFlow, PyTorch и JAX, для эффективного обучения и исполнения на разных графических процессорах и специализированных ускорителях. К совместной работе над проектом присоединились такие компании, как Google, NVIDIA, AMD, Intel, Meta, Apple, Arm, Alibaba и Amazon.
Проект OpenXLA предоставляет современный компилятор машинного обучения, который может масштабироваться в условиях сложности инфраструктуры машинного обучения. Его фундаментальными столпами являются производительность, масштабируемость, переносимость, гибкость и расширяемость для пользователей. С OpenXLA мы стремимся раскрыть реальный потенциал ИИ, ускорив его разработку и внедрение.
OpenXLA позволяет разработчикам компилировать и оптимизировать модели из всех ведущих платформ машинного обучения для эффективного обучения и обслуживания на самых разных аппаратных средствах. Разработчики, использующие OpenXLA, увидят значительные улучшения во времени обучения, производительности, задержке обслуживания и, в конечном итоге, во времени выхода на рынок и затратах на вычисления.
Надеемся, что совместными усилиями основных исследовательских групп и представителей сообщества, можно будет стимулировать развитие систем машинного обучения и решить проблему фрагментации инфраструктуры для различных фреймворков и команд.
OpenXLA позволяет реализовать эффективную поддержку различного оборудования, независимо от фреймворка, на котором основана модель машинного обучения. Ожидается, что OpenXLA сократит время обучения модели, повысит производительность, сократит задержки, уменьшит вычислительные затраты и сократит время выхода на рынок.
OpenXLA состоит из трех основных компонентов, код которого распространяется под лицензией Apache 2.0:
- XLA (ускоренная линейная алгебра) — компилятор, позволяющий оптимизировать модели машинного обучения для высокопроизводительного выполнения на разных аппаратных платформах, включая GPU, CPU и специализированные ускорители различных производителей.
- StableHLO — это базовая спецификация и реализация набора высокоуровневых операций (HLO) для использования в моделях систем машинного обучения. Он действует как слой между платформами машинного обучения и компиляторами, которые преобразуют модель для работы на определенном оборудовании. Слои подготовлены для создания моделей в формате StableHLO для фреймворков PyTorch, TensorFlow и JAX. Пакет MHLO используется в качестве основы для StableHLO, который расширен за счет поддержки сериализации и контроля версий.
- IREE (Intermediate Representation Execution Environment) — это компилятор и среда выполнения, которые преобразуют модели машинного обучения в универсальное промежуточное представление на основе формата MLIR (Intermediate Multi-Level Representation) проекта LLVM. Из особенностей выделяются возможность прекомпиляции (досрочно), поддержка управления потоком, возможность использования динамических элементов в моделях, оптимизация под разные CPU и GPU, низкие накладные расходы.
Относительно основных преимуществ OpenXLA упоминается, что оптимальная производительность была достигнута без необходимости углубляться в написание кода зависит от устройства, в дополнение к обеспечить нестандартную оптимизацию, включая упрощение алгебраических выражений, эффективное распределение памяти, планирование выполнения с учетом сокращения максимального потребления памяти и накладных расходов.
Еще одним преимуществом является упрощение масштабирования и распараллеливание вычислений. Разработчику достаточно добавить аннотации для подмножества критических тензоров, на основе которых компилятор сможет автоматически сгенерировать код для параллельных вычислений.
Также подчеркивается, что мобильность обеспечивается поддержкой нескольких аппаратных платформ, такие как графические процессоры AMD и NVIDIA, процессоры x86 и ARM, ускорители машинного обучения Google TPU, IPU AWS Trainium Inferentia, Graphcore и Wafer-Scale Engine Cerebras.
Поддержка подключения расширений с реализацией дополнительных функций, как поддержка написания примитивов глубокого машинного обучения с использованием CUDA, HIP, SYCL, Triton и других языков для параллельных вычислений, а также возможность ручной корректировки узких мест в моделях.
Наконец, если вам интересно узнать об этом больше, вы можете проконсультироваться с подробности по следующей ссылке.