Inferência rápida e portátil do Llama2 em edge heterogênea

Stack Rust+Wasm: uma alternativa poderosa ao Python

• Aplicações Rust+Wasm têm tamanho 1/100 do Python, são 100 vezes mais rápidas e podem ser executadas com segurança em qualquer lugar sem qualquer alteração na aceleração de hardware.
• Rust é a linguagem da AGI (inteligência artificial geral), e foi criado um programa simples em Rust para fazer inferência do modelo llama2 em velocidade nativa.
• A aplicação binária compilada para Wasm (2 MB) é totalmente portável entre dispositivos com diferentes aceleradores de hardware.
• O runtime Wasm (WasmEdge) oferece um ambiente de execução seguro e protegido na nuvem.

Instalação e execução do programa Rust

• Step 1. Instalar o WasmEdge e o plugin GGML: em dispositivos Linux ou Mac, a instalação é feita com os comandos abaixo.
• Step 2. Baixar o app Wasm pré-compilado e o modelo: baixe o modelo llama2 no formato GGUF; no exemplo, é usada a versão do modelo llama2 7B chat ajustada e quantizada com pesos de 5 bits.
• Step 3. Executar: rode a aplicação de inferência wasm com o WasmEdge, passe o modelo GGUF e insira perguntas para conversar com o modelo.

Configuração do funcionamento do modelo

• É possível configurar a forma de interação com o modelo usando opções de linha de comando.
• Por exemplo, é possível definir um comprimento de contexto específico, o número máximo de tokens a prever em cada resposta, imprimir estatísticas e transmitir a resposta do modelo para stdout.
• O programa gera cerca de 25 tokens por segundo em um MacBook M2 básico.

Agentes e apps de LLM

• Um servidor de API compatível com OpenAI é criado usando Rust e WasmEdge.
• Com isso, é possível criar agentes e apps de LLM usando ferramentas de desenvolvimento compatíveis com OpenAI.

Vantagens de Rust+Wasm em vez de Python

• A stack Rust+Wasm é uma alternativa poderosa à stack Python para aplicações de inferência de IA, oferecendo infraestrutura de computação em nuvem.
• Ultraleve: a aplicação de inferência tem apenas 2 MB, incluindo todas as dependências.
• Muito rápida: oferece velocidade nativa de C/Rust em todas as partes da aplicação de inferência.
• Portabilidade: a mesma aplicação em bytecode Wasm pode ser executada em todas as principais plataformas de computação que suportam aceleração de hardware heterogênea.
• Fácil de instalar, desenvolver e implantar: sem dependências complexas, é possível compilar um único arquivo Wasm com ferramentas padrão em um notebook e implantá-lo em qualquer lugar.
• Segura e pronta para a nuvem: o runtime Wasm foi projetado para isolar código de usuário não confiável, pode ser gerenciado por ferramentas de contêiner e implantado facilmente em plataformas cloud-native.

Programa de inferência em Rust

• O programa de inferência de demonstração foi escrito em Rust e compilado para Wasm.
• O código-fonte em Rust é muito simples, composto por apenas 40 linhas de código.
• O programa em Rust gerencia a entrada do usuário, acompanha o histórico da conversa, converte o texto para o template de chat do llama2 e executa a tarefa de inferência usando a API WASI NN.

Execução na nuvem ou na edge

• Depois de ter o arquivo de bytecode Wasm, é possível implantá-lo em qualquer dispositivo que suporte o runtime WasmEdge.
• O WasmEdge é instalado com o plugin GGML e usa automaticamente diversas acelerações de hardware para executar o modelo llama2.

Próximos passos

• As ferramentas GGML do WasmEdge já estão disponíveis e vêm sendo usadas por clientes cloud-native, mas ainda estão em estágio inicial.
• Se você tem interesse em contribuir com o projeto open source e ajudar a moldar a direção da futura infraestrutura de inferência para LLMs, há algumas tarefas de baixo nível nas quais é possível contribuir.

Outros modelos de IA

• WasmEdge e WASI NN permitem construir aplicações de inferência que, como alternativa ao Python, são leves, rápidas, portáveis e seguras.
• Por exemplo, o projeto mediapipe-rs fornece uma API Rust+Wasm para os modelos Tensorflow mediapipe do Google.

Opinião do GN⁺

O ponto mais importante deste texto é apresentar Rust e WebAssembly (Wasm) como alternativa ao Python para inferência de IA, destacando vantagens significativas em tamanho, velocidade, portabilidade e segurança. Essa stack tecnológica integra infraestrutura de computação em nuvem para melhorar o desempenho em diversos dispositivos e ambientes de nuvem, resolver problemas complexos de dependência e simplificar o desenvolvimento e a implantação. Essa abordagem pode ter um impacto importante especialmente no futuro da computação de edge e da inferência de IA, tornando o tema bastante interessante para quem acompanha tecnologia.

[GN⁺]

Comentários do Hacker News

• Embora demonstre apoio a Rust e WASM, o código real não passa de um script básico de linha de comando escrito em Rust, e o processamento principal depende de uma única linha de código que chama o backend WASI-NN. Neste caso, ele é fornecido pelo runtime WasmEdge, que é escrito em C++.
• Para desenvolvedores de aplicações, este projeto oferece a forma mais fácil de empacotar o motor de inferência como um arquivo distribuível. Os pesos do modelo já são portáveis e podem ser baixados conforme necessário.
• Levanta a dúvida se o projeto não seria apenas um wrapper para llama.cpp e compartilha experiência no desenvolvimento de frameworks de ML com Rust e WebGPU.
• Aponta que um projeto dependente de wasm-nn na prática não é portátil, e pode ser impossível executá-lo caso não haja suporte ao hardware específico. Isso entra em conflito com o propósito do WASI.
• Questiona a precisão da informação de que o build para Mac OS do plugin GGML executa inferência usando a API Metal na GPU, e não no Neural Engine dos M1/M2/M3.
• Demonstra preocupação de que, embora possam surgir formas de executar a arquitetura Transformer de maneira mais eficiente e rápida na edge computing, a exigência de VRAM continuará sendo um gargalo importante, podendo limitar o carregamento de modelos maiores.
• Pergunta se o projeto pode rodar offline em um iPhone. Isso pode ser útil em situações em que não há acesso à internet.
• Critica o marketing de projetos que usam wrappers escritos em outras linguagens como se fossem inovadores em tamanho ou velocidade. Na prática, trata-se de um projeto baseado em llama.cpp, e a velocidade de inferência ou o tamanho do binário não são limitados pela escolha do wrapper.
• Pergunta como reduzir o tamanho sem perder precisão quando os pesos treinados do modelo ocupam uma parte significativa do volume total.
• O núcleo do código-fonte em Rust é bastante simples, consistindo em 40 linhas que gerenciam a entrada do usuário, acompanham o histórico da conversa, convertem o texto para o template de chat do llama2 e executam a inferência usando a API WASI NN.