Introdução ao MLX para Apple Silicon

Olá! Meu nome é Awni. Hoje, tenho o prazer de apresentar o MLX. O MLX é um framework de matriz de código aberto desenvolvido especificamente para o Apple Silicon. É uma ferramenta flexível para uso com cálculos numéricos básicos até a execução de modelos de aprendizado de máquina de grande escala em dispositivos Apple. Se você quer gerar textos com grandes modelos de linguagem ou gerar imagens, áudios ou até vídeos com os modelos mais recentes, o MLX é ideal. Você também pode usá-lo para treinar, ajustar ou personalizar modelos de aprendizado de máquina no Mac. Vou começar contando um pouco mais sobre o que é o MLX e onde ele é útil. Também abordarei os fundamentos de uso do MLX em Python, como instalação e operações básicas de matriz. Depois, falarei sobre alguns dos principais recursos que diferenciam o MLX de outros frameworks. A seguir, apresentarei algumas ferramentas que o MLX tem para agilizar cargas de trabalho de aprendizado de máquina no Apple Silicon.

Na última seção, apresentarei brevemente o MLX Swift e mostrarei como começar a usá-lo. Vamos começar com a introdução ao MLX. O MLX foi criado do zero para ser rápido no Apple Silicon, onde pode ser executado na CPU ou acelerado na GPU. Podemos usar o MLX em vários apps, desde computação numérica em pequena escala até aprendizado de máquina em grande escala. Ele foi criado para ser prático e flexível, sem comprometer a velocidade e a eficiência.

O MLX tem uma API principal parecida com o NumPy. Muitas vezes, você pode usá-la como substituto drop-in para acelerar vários cálculos numéricos. O MLX também tem as ferramentas necessárias para aprendizado de máquina, incluindo diferenciação automática e bibliotecas de nível superior. As APIs de alto nível são parecidas com o PyTorch e o JAX. Se você conhece qualquer um desses frameworks, o MLX será familiar e ainda mais fácil de usar. Você pode usar o MLX para aprendizado de máquina avançado diretamente no dispositivo. Por exemplo, o MLX é usado pelo LM Studio, um app popular para gerar texto com grandes modelos de linguagem no Mac.

Você pode usar o pacote MLX LM, criado com base no MLX, para gerar texto e ajustar modelos de linguagem de até centenas de bilhões de parâmetros em tamanho. Para saber mais sobre isso, confira a sessão "Explorar grandes modelos de linguagem no Apple Silicon com o MLX". O MLX tem uma API Python completa, que é útil para prototipagem rápida. Ele também tem uma API no Swift, que inclui todos os pacotes de alto nível para criar e otimizar redes neurais. O MLX também tem APIs em C++ e C. Use o MLX em uma dessas linguagens para rodar os modelos de aprendizado de máquina mais recentes no Apple Silicon, como Mac, iPhone, iPad e Vision Pro. Todo o software MLX tem código aberto de acordo com uma licença MIT permissiva. O software principal está disponível no GitHub, bem como vários exemplos e pacotes criados com as APIs Python e Swift. O MLX também tem uma comunidade ativa de criadores de modelos no Hugging Face. Vários modelos recentes já estão na organização Hugging Face da comunidade do MLX, e novos modelos são carregados todos os dias. O jeito mais fácil de começar a usar o MLX em Python é instalar a partir do PyPi. Basta executar uma linha no terminal, pip3 install mlx. O MLX é prático. Para começar a executar cálculos em matrizes, basta abrir um arquivo Python e importar o MLX. Depois, você pode criar algumas matrizes e fazer operações básicas nelas. Por exemplo, aqui vamos adicionar duas matrizes inteiras.

Você pode inspecionar facilmente informações sobre uma matriz, como forma e tipo de dados. Como disse antes, a API Python do MLX é semelhante ao NumPy. As operações costumam ter os mesmos nomes e assinaturas e se comportar da mesma maneira. Se você conhece o NumPy ou um framework semelhante, o MLX será familiar e fácil de usar. Agora que conhece o que é MLX, onde ele é útil e alguns conceitos básicos, vamos conhecer os principais recursos. Isso inclui memória unificada, avaliação lenta, transformações de função e alguns pacotes de nível superior para criar e otimizar redes neurais. O MLX foi projetado para aproveitar ao máximo o Apple Silicon. Isso inclui um novo modelo de programação específico para memória unificada.

A maioria dos sistemas comumente usados para aprendizado de máquina tem uma GPU dedicada com memória separada. O Apple Silicon, por outro lado, tem uma arquitetura de memória unificada. Isso significa que a CPU e a GPU compartilham a mesma memória física.

Para trabalhar com memória unificada, o MLX é diferente do que você pode conhecer em frameworks tradicionais. Em frameworks tradicionais, a computação segue dados. Se a matriz estiver na memória da CPU, a computação ocorrerá na CPU. Se a matriz estiver na memória da GPU, a computação ocorrerá na GPU.

No MLX, as matrizes são alocadas na memória unificada. Você nunca precisa copiar em outro lugar para usá-las em um dos dispositivos compatíveis.

Em vez disso, para executar uma operação em um dispositivo, especifique o dispositivo para a operação. Por exemplo, aqui vamos adicionar a e b na GPU e multiplicar a e b na CPU. Essas operações podem ser executadas em paralelo, e o MLX gerenciará automaticamente as dependências, quando houver. Outro recurso importante do MLX é a computação lenta. Para tornar o MLX o mais eficiente possível, especialmente para cálculos grandes, ele tem um mecanismo de execução lenta.

Quando uma operação, como adição em duas matrizes, é chamada, nenhuma computação real acontece. Em vez disso, um gráfico de computação é criado, como o visto aqui.

A matriz C ainda não foi computada. Ela só é computada se você realmente precisar usá-la. Por exemplo, imprimir C ou converter C do MLX para uma lista Python forçará o cálculo. Você pode forçar explicitamente o gráfico a ser avaliado usando mx.eval. A computação lenta tem vários benefícios. Ao dissociar a criação e a execução do gráfico de computação, o MLX pode fazer transformações e otimizações no gráfico antes de calcular os resultados. Além disso, com a computação lenta, você só paga pelo que usa. As transformações de função são outro recurso fundamental do MLX. Elas o elevam de um framework de matriz rápida para uma ferramenta mais avançada para treinar e melhorar o desempenho de modelos de aprendizado de máquina.

As transformações de função são aquelas que tomam funções como entrada e retornam novas funções como saída.

O MLX tem várias transformações de função. Elas normalmente se enquadram em uma de duas categorias. As transformações para diferenciação automática servem para otimizar o gráfico de computação. Por exemplo, você pode usar uma transformação de função para calcular automaticamente o gradiente de funções no MLX.

Imagine que você tenha uma função básica que calcula o seno da sua entrada. Para obter o gradiente dessa função, você pode usar a transformação de função mx.grad. A saída de mx.grad é uma nova função, que, quando invocada em uma matriz, fornece a derivada. As transformações de função são arbitrariamente componíveis. Você pode obter a segunda derivada simplesmente usando mx.grad na saída de mx.grad.

O resultado é outra função, que, quando chamada em uma matriz, apresenta a segunda derivada do sinal.

O MLX também tem dois pacotes de nível superior para facilitar a criação e o treinamento de redes neurais.

O primeiro é mlx.nn, uma biblioteca modular usada para construir redes neurais. O segundo é mlx.optimizers, uma biblioteca de algoritmos de otimização comuns. Os dois pacotes podem ser usados de forma independente, mas também se integram perfeitamente um ao outro.

O pacote mlx.nn tem todas as funcionalidades necessárias para criar redes neurais. O nn.module é a classe básica principal, da qual todas as camadas e contêineres herdam. Ele expõe métodos úteis para acessar, carregar, salvar parâmetros e muito mais.

A biblioteca nn também tem um conjunto de camadas padrão prontas, como nn.Linear, mas é possível criar suas próprias camadas herdando de nn.Module.

As funções de perda e os métodos de inicialização mais usados também estão incluídos nos subpacotes nn.losses e nn.init.

Vamos conferir como você pode criar uma rede neural multicamada simples com mlx.nn.

A primeira etapa é criar uma classe personalizada que herda de nn.Module. Nesse caso, usaremos uma rede neural de camada oculta simples. Criamos as camadas lineares dentro do método de inicialização do módulo. Em seguida, implementamos a função call, que calcula a saída do módulo conforme a entrada fornecida. A função de chamada chama a primeira camada linear, aplica a função de ativação relu e, em seguida, chama a segunda camada linear. Embora o MLX seja projetado e otimizado para a arquitetura de memória unificada do Apple Silicon, para facilitar seu uso, o pacote de rede neural de nível superior foi criado para ser semelhante aos frameworks de aprendizado de máquina mais usados, como o PyTorch.

Vamos comparar a implementação do modelo MLX com a mesma implementação no PyTorch. Elas são quase idênticas, com apenas duas pequenas diferenças na função que calcula a saída. Se você já criou modelos no PyTorch antes, mudar para o MLX deve ser muito simples. Agora que conhece a maioria dos principais recursos do MLX, vamos ver como usá-lo para aprimorar suas cargas de trabalho de aprendizado de máquina. Vou começar mostrando como compilar funções para acelerá-las. Depois, falarei sobre o subpacote mx.fast, que tem implementações rápidas de operações comuns de aprendizado de máquina e uma API para adicionar kernels do Metal personalizados. Em seguida, mostrarei como usar a quantização para reduzir o uso de memória e acelerar a execução do modelo. Por fim, mostrarei como usar o MLX para distribuir uma computação em muitas máquinas. Quase toda computação realista do MLX consistirá em funções que executam várias operações em matrizes. Uma maneira simples de agilizar funções como essa é com a transformação de função mx.compile. Imagine que você tenha uma função que faz algumas operações em termos de elemento, como a função de ativação GELU mostrada aqui. O gráfico de computação para essa função tem vários nós nele. Cada um desses nós corresponde a uma inicialização separada do Kernel da GPU nos bastidores.

A compilação do gráfico usa todos esses Kernels separados em um único Kernel combinado. Isso pode reduzir o consumo de largura de banda de memória e a sobrecarga gráfica, resultando em cálculos mais eficientes.

Usar mx.compile é tão fácil quanto decorar a função com a transformação de função mx.compile. A compilação costuma funcionar bem, mas para operações complexas, incluindo operações comuns no aprendizado de máquina, é provável que elas sejam executadas mais rápido usando o subpacote mx.fast. Por exemplo, muitos dos principais fundamentos de um modelo de transformador usam operações em mx.fast. Isso inclui codificações posicionais, camadas de normalização e atenção do produto de pontos de escala. As operações no mx.fast são mais especializadas, mas altamente ajustadas para serem o mais rápidas possível para treinamento e inferência. Elas também são altamente configuráveis, de modo que podem acelerar muitas variações de uma dada computação. Por exemplo, a operação de atenção ao produto de ponto de escala pode usar uma máscara opcional como entrada. A máscara pode ser uma máscara aditiva, uma máscara booleana ou uma cadeia de caracteres indicando o tipo de máscara. Vamos ver mais detalhes sobre a norma RMS, que é uma das operações no mx.fast. A norma RMS é usada em quase todos os grandes modelos de linguagem modernos baseados em transformadores. Uma implementação básica usando operações MLX resulta em um grande gráfico de computação.

Em vez disso, podemos substituir toda a implementação por uma única chamada para mx.fast.rms_norm. O código é mais simples, o gráfico de computação tem um único nó e a computação em si será muito mais rápida.

O MLX tem uma API para adicionar Metal Kernels personalizados para casos em que sua função poderia se beneficiar de uma implementação mais personalizada, e ela ainda não está no mx.fast. Você escreve o Kernel personalizado, e o MLX lida com todo o resto, incluindo compilação e execução just-in-time. Esses Kernels são escritos em Metal, que é uma linguagem e API desenvolvida pela Apple para executar funções em GPUs Apple. Você cria o Kernel enviando uma string do código-fonte do Metal, bem como algumas informações sobre as entradas e saídas.

Você chama o Kernel especificando o tamanho da grade e as formas e os tipos da saída. O MLX trata a chamada do Kernel da mesma forma que qualquer outra operação. Ele cria um nó no gráfico de computação e é avaliado lentamente.

Outra ferramenta no nosso kit de ferramentas para agilizar suas cargas de trabalho de aprendizado de máquina é a quantização. Os modelos grandes precisam de muita memória e muita largura de banda da memória para serem rápidos. Em muitos casos, a precisão exigida para o treinamento é muito maior do que o realmente necessário para inferência para obter uma qualidade igual ou similar. Ao reduzir a precisão, você pode utilizar modelos maiores na memória e executá-los mais rapidamente.

Se o seu modelo tiver precisão de ponto flutuante de 32 bits, use bfloat16 ou float16 para reduzir os requisitos de memória pela metade. Quando 16 bits é muito, o MLX tem rotinas integradas para quantizar matrizes para serem ainda menores e fazer operações com elas. Por exemplo, você pode quantizar usando 4 bits por elemento para reduzir ainda mais os requisitos de memória.

Para quantizar uma matriz, use mx.quantize. Você especifica o número de bits a serem usados por elemento e o tamanho do grupo. O tamanho do grupo determina o número de elementos na matriz quantizada, que obtém uma escala compartilhada e um valor de viés. Quanto menores os bits e maior o tamanho do grupo, menor e mais rápido o resultado. O MLX tem várias opções de bits e tamanhos de grupo para oferecer o máximo de flexibilidade possível.

Você pode multiplicar qualquer vetor ou matriz de entrada não quantizado por sua matriz quantizada usando mx.quantized_matmul.

Use mx.dequantize para recuperar uma aproximação à entrada original.

O MLX.mn também tem um utilitário prático para quantizar um módulo em um único comando. Digamos que você tenha um modelo que é uma pilha de uma camada de integração seguida por várias camadas lineares. Você pode quantizar o modelo com nn.quantize. O comando quantize usa um callback opcional para permitir um controle mais refinado sobre quais camadas quantizar e qual precisão usar para uma camada.

Ao gerar texto com grandes modelos de linguagem, a quantização reduz consideravelmente o uso de memória e aumenta os tokens gerados por segundo. Em alguns casos, uma máquina simplesmente não é suficiente. Por exemplo, talvez você queira gerar texto com um modelo grande que não cabe na memória de uma única máquina. Ou você talvez queira ajustar ou avaliar um modelo em um grande conjunto de dados. As duas funções são fáceis de paralelizar e podem ser mais rápidas com várias máquinas.

O MLX possui a capacidade de executar cálculos arbitrários em várias máquinas. As máquinas podem ser conectadas via ethernet ou Thunderbolt.

Use o subpacote mx.distributed para distribuir cálculos em várias máquinas. O MX distribuído é basicamente um conjunto de operações de comunicação.

Por exemplo, all_Sum adiciona as matrizes de entrada em todas as máquinas. A saída de all_Sum é a entrada somada, que é a mesma para cada máquina. Vamos ver mais detalhes de como somar uma matriz em várias máquinas. Inicialize o back-end distribuído usando mx.distributed.init. Essa etapa é opcional, mas é o que você chama se precisar acessar o grupo de comunicação. O grupo de comunicação tem informações úteis anexadas a ele, como o número total de processos e o índice de processos atual.

Em seguida, crie uma matriz com um único valor em cada processo e chame mx.distributed.all_sum para somar as matrizes em todos os processos.

O MLX tem um inicializador prático para executar um programa do MLX em várias máquinas. Para executar um programa em 4 máquinas, use mlx.launch com os 4 endereços IP do host. Tudo o que mencionei até agora foi usando o MLX em Python. Em muitos casos, você pode preferir a facilidade e a flexibilidade do Python. Em outros casos, você pode preferir o Swift. Por esse motivo, o MLX tem uma API completa no Swift. Ele foi criado sobre o Metal e funciona muito bem em macOS, iOS, iPadOS, visionOS e muitos outros. Começar a usar o MLX no Swift é tão fácil quanto adicioná-lo como um pacote ao seu projeto Xcode. Clique no projeto e, em seguida, clique no sinal de adição na guia de dependências do pacote. Em seguida, insira o URL do repositório MLX Swift no GitHub e clique no botão Adicionar pacote. Basta isso para começar a criar com o MLX Swift. Para facilitar ao máximo a transição entre Python e Swift, as APIs entre as duas linguagens são intencionalmente semelhantes. Aqui está uma comparação lado a lado de um trecho de código Python, que vimos antes, com o equivalente no MLX Swift. Criar matrizes, executar operações nelas e inspecionar metadados sobre a matriz é quase igual no Swift e no Python. Todos os principais recursos que vimos usando MLX no Python, bem como as otimizações que discutimos na seção Acelerar o MLX, também se aplicam ao MLX Swift. Tivemos uma visão geral ampla de vários dos principais recursos do MLX. Para saber mais sobre o framework, confira o site do MLX, que tem links para a documentação, exemplos e muito mais. As APIs Python e Swift têm um repositório de exemplos, com vários casos de uso de aprendizado de máquina, incluindo treinamento e geração de modelos de linguagem, geração de imagens, reconhecimento de fala e muito mais. Esses exemplos são um ótimo jeito de aprender mais sobre o MLX e bons pontos de partida para criar com MLX em seu projeto. Agradeço a participação. Mal posso esperar para ver o que você vai criar com o MLX.