MLX

MLX é um framework de arrays para aprendizado de máquina em Apple silicon que também funciona com CUDA. Em Apple silicon, os arrays permanecem na memória compartilhada para evitar cópias de dados entre a CPU e a GPU. A computação preguiçosa (lazy computation) permite a manipulação de grafos e otimizações. O suporte nativo a safetensors significa que os modelos de linguagem do Transformers rodam diretamente no MLX.

Instale a biblioteca mlx-lm.

pip install mlx-lm transformers

Carregue qualquer modelo de linguagem do Transformers do Hub, desde que a arquitetura do modelo seja suportada. Não é necessária conversão de pesos.

from mlx_lm import load, generate

model, tokenizer = load("openai/gpt-oss-20b")
output = generate(
    model,
    tokenizer,
    prompt="The capital of France is",
    max_tokens=100,
)
print(output)

Integração com Transformers

• mlx_lm.load carrega os pesos em safetensor e retorna um modelo e um tokenizer.
• O MLX carrega arrays de pesos indexados pelos nomes dos tensores e os mapeia para uma árvore de parâmetros de um nn.Module do MLX. Isso corresponde à forma como os checkpoints do Transformers são organizados.

A integração do MLX Transformers é bidirecional. O Transformers também pode carregar e executar pesos do MLX a partir do Hub.

Recursos

• Documentação do MLX
• repositório mlx-lm contendo implementações de LLMs em MLX
• Biblioteca da comunidade mlx-vlm com implementações de VLM