À medida que os LLMs se tornam mais capazes, as expectativas do usuário de velocidade e confiabilidade continuam aumentando, especialmente para aplicações de empresas e produtividade. É por isso que acreditamos no poder da IA no dispositivo, que pode oferecer experiências em tempo actual e de alta qualidade que combinam ou excedem soluções baseadas em nuvem ou híbridas. O desafio? Não existe um plano estabelecido para otimizar e escalar modelos no dispositivo.
Abordamos esse desafio de frente, desenvolvendo uma abordagem sistemática para otimizar os modelos de linguagem (como o T5) para executar o dispositivo-sem qualquer degradação da qualidade. Especificamente, reduzimos a memória e a latência do nosso modelo de correção de erro gramatical (GEC) em mais de 50%, permitindo que ele seja executado com eficiência nos dispositivos dos usuários. Também criamos um package de desenvolvimento de software program personalizado (SDK) para fornecer esse modelo otimizado a milhões de usuários em várias plataformas de desktop. Ao enviar com sucesso essa funcionalidade em escala, demonstramos que nossa abordagem no dispositivo funciona, tornando-nos uma das empresas pioneiras nesse espaço.
Nesta postagem do weblog, compartilharemos como resolvemos os desafios técnicos de otimizar e escalar nosso modelo GEC, estabelecendo uma base para o desenvolvimento futuro da IA no dispositivo.
O que é preciso para executar o modelo GEC no dispositivo?
Embora a Grammarly forneça mais do que apenas a correção de erros gramaticais, é uma parte central de como capacitamos os usuários a se comunicar de maneira eficaz. Para garantir uma experiência de escrita perfeita, a gramática deve fornecer sugestões de alta qualidade em tempo real-menos de 100 milissegundos. Conseguir isso requer resolver três desafios únicos:
- Gerenciamento de memória: os dispositivos de usuário têm memória limitada, que geralmente é compartilhada com outros aplicativos, tornando a otimização do uso da memória crítico para o desempenho. Nosso modelo unique (projetado para servidores em nuvem) precisava de quase 4 GB, aproximadamente o tamanho da RAM de uma área de trabalho médio – sendo impraticável para executar localmente sem otimização significativa.
- Eficiência computacional: como na memória, os dispositivos de usuário têm poder de processamento limitado que geralmente é compartilhado com outros aplicativos. Para uma experiência em tempo actual, as operações com recursos de recursos do modelo (como a inferência) devem funcionar rapidamente sem comprometer a qualidade. Se nosso modelo exigir muitos recursos, ele pode causar atraso, interferir em outros aplicativos ou drenar rapidamente a duração da bateria, levando a uma má experiência do usuário.
- Implantação de plataforma cruzada: para atender a todos os usuários gramaticais, nossa solução deve funcionar em diferentes plataformas e dispositivos. Isso é desafiador, porque cada plataforma usa diferentes {hardware}, linguagens de programação e APIs de aprendizado de máquina. Também devemos considerar os recursos variados do dispositivo – de poderosos MacBooks poderosos com GPUs dedicadas ao orçamento do Chromebooks com recursos mínimos.
Para superar esses desafios, abordamos iterativamente cada problema: gerenciamento de memória, latência computacional e implantação de plataformas cruzadas.
Reduzindo a pegada de memória
Para garantir que o GEC não afetaria o uso da memória de outros aplicativos, definimos um alvo ambicioso de otimizar o modelo para ser executado abaixo de 1 GB de memória. Uma técnica comum para resolver esse problema é a quantização, na qual os pesos do modelo (normalmente carros alegóricos de 32 bits) são convertidos em números menores e menos precisos (como um número inteiro de 4 bits). Embora a quantização reduz bastante a pegada de memória, ela também pode degradar a precisão do modelo.
Para equilibrar a qualidade do modelo e a otimização da memória, experimentamos diferentes níveis de quantização, incluindo BFLOAT16 e INT4. O BFLOAT16 teve os melhores resultados: teve um impacto mínimo na precisão ao reduzir o uso da memória em 50%. Ao combinar quantização com outras técnicas de otimização (como as otimizações de gráficos descritas abaixo), o modelo last de correção de erro gramatical foi executado em menos de 300 MB.
Melhorando a eficiência computacional
Para garantir uma experiência do usuário em tempo actual, calculamos que todo o modelo GEC deve ser executado em mais de 100 tokens/segundo, o que exigiu que reduzissem a latência em 50%.
Começamos otimizando o modelo T5, o mecanismo de inferência do núcleo para o modelo GEC. Acreditávamos que aumentar a velocidade do modelo T5 de 70 tokens/segundo a 200 tokens/segundo nos permitiria aplicar esses aprendizados a outros componentes do pipeline e cumprir nossas metas gerais de desempenho.
Então, começamos realizando otimizações de gráficos manuais no modelo, reorganizando e simplificando operações de maneiras semelhantes às realizadas por bibliotecas como Tensorrt e Onnx Runtime. Em seguida, examinamos os cálculos dentro de cada operação. Algumas otimizações foram óbvias, como a remoção de operações desnecessárias (por exemplo, tipo de tipo de tipo desnecessário), minimizando operações intensivas em tempo (por exemplo, remodelam e transpomue operações) e operações de sequenciamento para manter os dados no mesmo processador o máximo possível (transferências de dados são intensivas no tempo).
Um grande desbloqueio foi utilizar um kernel otimizado e operações fundidas que alavancam a especificação de {hardware} do dispositivo. Por exemplo, um cálculo crítico em uma rede neural é a operação de atenção de várias cabeças (MHA), que ocorre quando o modelo examina a entrada do usuário para determinar a consulta e a saída potencial do usuário. Esta operação requer vários cálculos simultaneamente em diferentes partes da consulta de entrada. Substituímos a implementação ingênua do MHA por uma função otimizada da MLX, a estrutura ML da Apple, acelerando assim o desempenho do modelo.
Com esses ganhos de desempenho arquitetônico e computacional, aumentamos nossa velocidade de processamento T5 de 70 tokens/segundo para 297 tokens/segundo – atendendo aos nossos objetivos de latência.
Implantação de plataforma cruzada sem costura
Agora que tínhamos um modelo de desempenho, enfrentamos um novo desafio: implantando -o para milhões de usuários de área de trabalho da Grammarly em várias plataformas, cada uma com seu {hardware}, idioma e APIs específicos para operações de ML. Uma abordagem direta é criar um SDK separado para cada plataforma. No entanto, essa abordagem faz com que a manutenção e iteração futura complicadas, pois cada alteração de código deve ser aplicada a cada SDK.
Para resolver isso, construímos um SDK baseado em ferrugem que executa modelos de aprendizado de máquina gramatical em três plataformas principais: Mac, Home windows e Extensão Chrome. Essa abordagem nos permite escrever um código uma vez e compilá -lo para cada plataforma, simplificando a implantação e a manutenção do modelo. Além disso, o SDK alavancou as bibliotecas ML da plataforma nativa (como os dispositivos Metallic for Apple), permitindo-nos aproveitar as acelerações específicas de {hardware}.
O futuro está no dispositivo
Usando o nosso SDK, enviamos nosso modelo GEC no dispositivo para milhões de usuários de mesa gramatical. Os primeiros resultados não mostram degradação em qualidade ou desempenho em comparação com o modelo anterior baseado em nuvem. Isso confirma nossa crença de que a IA no dispositivo pode oferecer experiências em tempo actual e de alta qualidade sem depender de servidores em nuvem-lançando o que antes period apenas uma visão na realidade.
Este trabalho é apenas o primeiro passo em direção a novas e poderosas experiências de IA. Aproveitando esses aprendizados de otimização e nosso SDK, estamos construindo versões no dispositivo de outros modelos complexos de escrita. Se você deseja enfrentar problemas desafiadores de otimização de IA em escala, venha trabalhar conosco. Confira o nosso Página de empregos.
Agradecimentos especiais a toda a equipe que trabalhou neste projeto: Sri Malireddi, Illia Dzivinskyi, Ignat Blazhko, Dhruv Matani e John Blatz.
