Adicionando suporte a CUDA ao MLX da Apple

• MLX é um framework de arrays para pesquisadores de aprendizado de máquina com Apple Silicon, oferecendo uma usabilidade semelhante à do NumPy e do PyTorch
• Está em andamento o trabalho para adicionar um backend CUDA ao MLX
• O objetivo principal é oferecer suporte a memória unificada (unified memory) e ampliar o uso em hardware NVIDIA
• No momento, apenas os exemplos dos tutoriais funcionam, e o build e os testes foram verificados em um ambiente Ubuntu 22.04 + CUDA 11.6
• O processo de otimização e refatoração está sendo repetido continuamente à medida que problemas iniciais de desempenho e gargalos são melhorados
• O backend CUDA do MLX está sendo desenvolvido com apoio da Apple, e o objetivo é oferecer uma experiência consistente ao desenvolver no Mac e implantar em grandes ambientes de computação, por meio de memória unificada e suporte a hardware NVIDIA

---

Introdução ao framework MLX

• MLX é um framework de arrays para aprendizado de máquina desenvolvido pela equipe de pesquisa em machine learning da Apple
• Foi criado para que engenheiros e pesquisadores possam experimentar de forma eficaz e validar ideias rapidamente
• Seu diferencial é o suporte ao Apple Silicon, mas agora, com a adição de um backend CUDA, busca expandir sua possibilidade de funcionamento também em ambientes com GPU da NVIDIA

Principais características

• API familiar: a API em Python é semelhante ao NumPy, e os pacotes de nível superior mlx.nn e mlx.optimizers oferecem uma interface parecida com a do PyTorch
  • APIs em C++, C e Swift também são fornecidas, garantindo usabilidade consistente com a API em Python
• Transformações de função componíveis: inclui recursos como diferenciação automática, vetorização automática e otimização de grafos computacionais
• Avaliação preguiçosa: estrutura que agrupa várias operações e calcula o resultado apenas quando necessário, permitindo uso mais eficiente de recursos
• Criação de grafo dinâmico: mesmo que o shape dos dados de entrada mude, a execução é imediata sem compilação lenta, o que facilita o desenvolvimento e a depuração
• Modelo multi-device e de memória unificada:
  • Antes, havia suporte apenas a CPU e GPU da Apple, mas com o backend CUDA está em andamento o suporte a ambientes com GPU da NVIDIA
  • Com a adoção de uma arquitetura de memória unificada, é possível operar sobre o mesmo objeto em todos os dispositivos compatíveis sem cópia de memória
• Design simples e extensível: permite que pesquisadores expandam e aprimorem o framework com facilidade

Diferenças em relação a outros frameworks e inspirações

• Recebe influência de design de NumPy, PyTorch, Jax, ArrayFire e outros
• Destaca especialmente o modelo de memória unificada, a interface simples e o ambiente de execução/depuração imediata

Exemplos representativos e áreas de uso

• O repositório MLX Examples contém diversos códigos práticos
  • Treinamento de modelos de linguagem Transformer
  • Geração de texto em larga escala com base em LLaMA e fine-tuning com LoRA
  • Geração de imagens com Stable Diffusion
  • Implementações de modelos recentes, como reconhecimento de fala baseado em OpenAI Whisper

O significado do MLX e o suporte a CUDA

• O MLX, que era um framework voltado ao Apple Silicon, está ampliando sua possibilidade de uso geral também em ambientes com GPU da NVIDIA com a introdução do backend CUDA
• Com suporte às versões mais recentes do CUDA, a perspectiva é que ele se consolide como um framework adequado tanto para pesquisa de ponta quanto para uso industrial, em hardware da Apple e da NVIDIA

[gosarinamu]

Era só permitir usar GPUs da NVIDIA no macOS... aff, aff.

[yangeok]

Boa notícia, hehe. Espero que adicionem suporte a CUDA logo, para que também seja possível fazer treinamento em alta velocidade no Mac~!

[GN⁺]

Comentários do Hacker News

• Quero confirmar se entendi corretamente: 1, programas compilados com MLX podem rodar em chips com suporte a CUDA. Mas 2, programas CUDA não podem rodar no Apple Silicon. Gostaria de saber se isso acontece porque o item 2 violaria direitos autorais, especialmente a famosa barreira de entrada da Nvidia
  • O item 1 está correto. Isso permite que desenvolvedores trabalhem em dispositivos Apple relativamente menos potentes, inclusive com UMA, e implantem o código final em sistemas Nvidia relativamente mais potentes. Isso é útil por vários motivos
  • O item 2 não é violação de direitos autorais. É possível reimplementar a API
  • Sobre o item 2, isso provavelmente não está certo mesmo. A HIP da AMD também está fazendo a mesma coisa desde que abandonou o OpenCL por volta de 2017–2018
  • Eu quero o item 3. Seria ótimo poder conectar uma GPU da NVIDIA ao Apple Silicon e usar CUDA. Imagino Apple Silicon, memória unificada, GPU e CUDA sendo usados juntos com PyTorch, JAX ou TensorFlow, mas ainda não tive uma experiência de verdade com MLX
  • O motivo principal de o item 2 não acontecer é que isso é bem mais difícil
• Se a sua reação foi "ué, mas a plataforma Apple não tem suporte oficial a CUDA?", então também vale notar que este conjunto de patches adiciona suporte até para "plataformas Linux com CUDA 12 e SM 7.0 (Volta) ou superior"
  Instruções de instalação
• Este projeto é liderado por zcbenz, conhecido por criar o Electron e outras coisas
  Sobre zcbenz
• Uma das principais características do MLX é o uso de arquitetura de memória unificada (UMA), e há curiosidade sobre como isso funciona
  Veja o bullet no readme do repositório: MLX no GitHub
  Parece interessante expandir isso para todas as APUs com UMA, independentemente do fabricante, mas fica a dúvida de como isso teria de funcionar de outra forma em GPUs discretas
  Lendo os comentários do PR, parece que o CUDA também dá suporte direto à API de UMA e faz as cópias de forma transparente quando necessário
  • Pela minha experiência, a falta de recursos de prefetch causa gargalos pesados de memória por causa da espera por cópias. Quando o conjunto inteiro de dados cabe na VRAM, dá para se virar com prefetch manual, mas na minha aplicação, treinamento de ML, a perda de desempenho era grande demais e acabei mudando diretamente para carregamento em streaming
• Mudando um pouco de assunto, muitos dos que trabalham no MLX parecem não estar oficialmente ligados à Apple. Como exemplo, veja o Twitter de prince_canuma
  Isso levanta a dúvida se a Apple está apoiando por trás dos panos sem se expor oficialmente. Também houve rumores de que a equipe estaria migrando
• A descrição do PR diz que é um "ongoing effort" para adicionar um backend CUDA
  Ou seja, isso significa que o código do MLX pode ser compilado e executado em hardware x86 + GeForce, e não o contrário
• Como disseram, "seria uma boa experiência de desenvolvimento escrever e testar no Mac e implantar em um supercomputador", então fica a dúvida se agora dá para usar MLX no Linux
  Testei diretamente e, no PyPI, só está publicada a versão para Python 3.12
  Página do MLX-CUDA no PyPI
• Fico pensando se isso tem relação com o Strix Halo. Sobrou orçamento no computador de trabalho e comprei um EVO-X2, que é parecido ou melhor do que máquinas EPYC znver5 de faixa intermediária. Supera com folga a maioria das instâncias EC2 ou GCE, até com NVMe. Esse hardware custa 1800
  Ultimamente também estão saindo reviews no YouTube dizendo que o DGX Spark tem pior custo-benefício do que o Strix Halo. A desvantagem é que o ROCm ainda é fraco, embora pareça que isso vá melhorar em breve. Se CUDA funcionasse perfeitamente em dispositivos Apple, seria uma opção a considerar seriamente, por mais barato e bom que o Strix seja
  • Strix Halo é basicamente o AMD Ryzen AI Max+ 395. Ele está sendo usado no Framework Desktop e em vários mini PCs. A largura de banda de memória desse chip é de 200GB/s, o que é muito bom para uma plataforma x86 comum, mas ainda fica muito atrás de GPUs Nvidia, por exemplo a 5090 com 1792GB/s, ou do Apple M3 Ultra com 800GB/s. O custo-benefício é excelente, mas para tarefas intensivas em memória, como LLMs, ainda parece estar só entrando no nível mínimo de desempenho aceitável
  • Pela descrição do PR, fica claro que o alvo é treinamento em clusters na nuvem
  • Já faz 10 anos que se fala que as operações de redes neurais em GPUs AMD logo competiriam com a Nvidia, e isso ainda não se concretizou
  • Fico curioso sobre como isso se compara ao Mac mini com M4
• Fiquei muito impressionado com os modelos do MLX. Dá para deixar modelos locais disponíveis para toda a família sem se preocupar com risco de incêndio como em um computador com Nvidia. Espero que o Apple Silicon se torne um concorrente sério no campo dos chips da Nvidia. Também fico me perguntando se esse suporte a CUDA pode ser uma estratégia de embrace, extend, extinguish (EEE)
• A Apple planeja, no futuro, construir datacenters baseados em chips da série M para usar em desenvolvimento de apps, testes e hospedagem de serviços externos