OpenXLA é um ecossistema compilador de ML de código aberto co-desenvolvido
Faz pouco, as maiores empresas envolvidas no desenvolvimento na área de aprendizado de máquina apresentadas o projeto OpenXLA, destinado ao desenvolvimento conjunto de ferramentas para compilar e otimizar modelos para sistemas de aprendizado de máquina.
O projeto se encarregou do desenvolvimento de ferramentas que permitem unificar a compilação de modelos elaborados nos frameworks TensorFlow, PyTorch e JAX para treinamento e execução eficientes em diferentes GPUs e aceleradores especializados. Empresas como Google, NVIDIA, AMD, Intel, Meta, Apple, Arm, Alibaba e Amazon se juntaram ao trabalho conjunto do projeto.
O Projeto OpenXLA fornece um compilador de ML de última geração que pode ser dimensionado em meio à complexidade da infraestrutura de ML. Seus pilares fundamentais são performance, escalabilidade, portabilidade, flexibilidade e extensibilidade para os usuários. Com o OpenXLA, aspiramos a liberar o potencial real da IA, acelerando seu desenvolvimento e entrega.
O OpenXLA permite que os desenvolvedores compilem e otimizem modelos de todas as principais estruturas de ML para treinamento e manutenção eficientes em uma ampla variedade de hardware. Os desenvolvedores que usam o OpenXLA verão melhorias significativas no tempo de treinamento, desempenho, latência do serviço e, finalmente, tempo de mercado e custos de computação.
Espera-se que unindo esforços das principais equipes de pesquisa e representantes da comunidade, será possível estimular o desenvolvimento de sistemas de aprendizado de máquina e resolver o problema de fragmentação da infraestrutura para várias estruturas e equipes.
O OpenXLA permite implementar suporte efetivo para vários hardwares, independentemente da estrutura na qual o modelo de aprendizado de máquina é baseado. Espera-se que o OpenXLA reduza o tempo de treinamento do modelo, melhore o desempenho, reduza a latência, reduza a sobrecarga de computação e reduza o tempo de lançamento no mercado.
OpenXLA consiste em três componentes principais, cujo código é distribuído sob a licença Apache 2.0:
- XLA (álgebra linear acelerada) é um compilador que permite otimizar modelos de aprendizado de máquina para execução de alto desempenho em diferentes plataformas de hardware, incluindo GPUs, CPUs e aceleradores especializados de vários fabricantes.
- StableHLO é uma especificação básica e implementação de um conjunto de operações de alto nível (HLOs) para uso em modelos de sistemas de aprendizado de máquina. Ele atua como uma camada entre os frameworks de aprendizado de máquina e os compiladores que transformam o modelo para rodar em hardware específico. As camadas são preparadas para gerar modelos no formato StableHLO para os frameworks PyTorch, TensorFlow e JAX. A suíte MHLO é usada como base para StableHLO, que é estendida com suporte para serialização e controle de versão.
- IREE (Intermediate Representation Execution Environment) é um compilador e tempo de execução que converte modelos de aprendizado de máquina em uma representação intermediária universal baseada no formato MLIR (Intermediate Multi-Level Representation) do projeto LLVM. Entre os recursos, destacam-se a possibilidade de pré-compilação (antecipadamente), suporte para controle de fluxo, capacidade de usar elementos dinâmicos em modelos, otimização para diferentes CPUs e GPUs e baixo overhead.
Em relação às principais vantagens do OpenXLA, é mencionado que o desempenho ideal foi alcançado sem ter que se aprofundar na escrita de código específico do dispositivo, além de fornecer otimizações prontas para uso, incluindo simplificação de expressões algébricas, alocação eficiente de memória, escalonamento de execução, levando em consideração a redução do consumo máximo de memória e sobrecargas.
Outra vantagem é a simplificação de escalonamento e paralelização de cálculos. Basta que um desenvolvedor adicione anotações para um subconjunto de tensores críticos, com base nos quais o compilador pode gerar automaticamente código para computação paralela.
Destaca-se também que a portabilidade é fornecida com suporte para várias plataformas de hardware, como GPUs AMD e NVIDIA, CPUs x86 e ARM, Google TPU ML Accelerators, AWS Trainium Inferentia IPUs, Graphcore e Wafer-Scale Engine Cerebras.
Suporte para conectar extensões com a implementação de funções adicionais, como suporte para escrever primitivas de aprendizagem de máquina profunda usando CUDA, HIP, SYCL, Triton e outras linguagens para computação paralela, bem como o possibilidade de ajuste manual de gargalos em modelos.
Por fim, se você estiver interessado em saber mais sobre o assunto, pode consultar o detalhes no link a seguir.