As maiores empresas que atuam no desenvolvimento em aprendizado de máquina apresentaram o projeto OpenXLA, destinado ao desenvolvimento colaborativo de ferramentas para compilação e otimização de modelos para sistemas de aprendizado de máquina. O projeto abrange o desenvolvimento de ferramentas que permitem unificar a compilação de modelos preparados nas estruturas TensorFlow, PyTorch e JAX, para treinamento e execução eficaz em diferentes GPUs e aceleradores especializados. Empresas como Google, NVIDIA, AMD, Intel, Meta, Apple, Arm, Alibaba e Amazon se juntaram ao projeto.
Espera-se que a colaboração de equipes de pesquisa líderes e representantes da comunidade estimule o desenvolvimento de sistemas de aprendizado de máquina e resolva o problema da fragmentação da infraestrutura para diferentes estruturas e equipamentos. O OpenXLA permite uma ampla compatibilidade de suporte a diferentes equipamentos, independentemente da estrutura em que o modelo de aprendizado de máquina foi criado. A expectativa é que, com o OpenXLA, o tempo de treinamento de modelos seja reduzido, a largura de banda aumentada, a latência diminuída, os custos de recursos computacionais reduzidos e o tempo de lançamento de produtos no mercado acelerado.
O OpenXLA é composto por três componentes principais, cujo código é distribuído sob a licença Apache 2.0:
- XLA (Accelerated Linear Algebra) — um compilador que permite otimizar modelos de aprendizado de máquina para execução de alto desempenho em várias plataformas de hardware, incluindo GPUs, CPUs e aceleradores especializados de diferentes fabricantes.
- StableHLO — uma especificação e implementação básica de um conjunto de operações de alto nível (HLO, High-Level Operations) para uso em modelos de sistemas de aprendizado de máquina. Serve como uma camada intermediária entre as estruturas de aprendizado de máquina e os compiladores que transformam o modelo para execução em equipamentos específicos. As camadas para geração de modelos no formato StableHLO foram preparadas para as estruturas PyTorch, TensorFlow e JAX. O conjunto MHLO foi usado como base para o StableHLO, que foi expandido com suporte para serialização e versionamento.
- IREE (Ambiente de Execução de Representação Intermediária) — um compilador e runtime que transforma modelos de aprendizado de máquina em uma representação intermediária universal, baseada no formato MLIR (Representação Intermediária de Múltiplos Níveis) do projeto LLVM. Entre suas características, destaca-se a possibilidade de compilação antecipada (ahead-of-time), suporte ao gerenciamento de fluxo, uso de elementos dinâmicos em modelos e otimização para diferentes CPUs e GPUs, com baixas sobrecargas.
Principais benefícios da ferramenta OpenXLA:
- Alcançar desempenho ideal sem a necessidade de se aprofundar na escrita de código específico para determinados dispositivos. Fornecimento de otimizações prontas, incluindo simplificação de expressões algébricas, alocação eficiente de memória e planejamento da execução com foco na redução do pico de consumo de memória e sobrecargas.
- Facilidade na escalabilidade e paralelização de cálculos. O desenvolvedor precisa apenas adicionar anotações para um subconjunto crítico de tensores, com base nos quais o compilador pode gerar automaticamente código para cálculos paralelos.
- Garantia de portabilidade por meio do suporte a várias plataformas de hardware, como GPUs AMD e NVIDIA, CPUs baseadas na arquitetura x86 e ARM, aceleradores de ML TPU do Google, IPUs da AWS Trainium Inferentia, Graphcore e Cerebras Wafer-Scale Engine.
- Suporte para a conexão de extensões com a implementação de recursos adicionais, como o suporte à escrita de primitivos de aprendizado de máquina profundo usando CUDA, HIP, SYCL, Triton e outras linguagens para cálculos paralelos. Possibilidade de ajuste manual de gargalos nos modelos.
Fonte: opennet.ru
