MaskLLM
NVIDIA Research · Programação · Outros
MaskLLM é um método de poda de código aberto da NVIDIA Research que aprende máscaras de esparsidade N:M para encolher modelos de linguagem e cortar o custo de inferência sem perder muita precisão.
Sobre MaskLLM
What Is MaskLLM
MaskLLM é um método de pesquisa para comprimir modelos de linguagem grandes. Ele faz parte da família de ferramentas de poda que removem os parâmetros dos quais o modelo não precisa, e faz isso por meio de esparsidade semiestruturada, também chamada de esparsidade N:M. A ideia é simples: em cada bloco de M pesos consecutivos, mantêm-se apenas N. Um padrão 2:4 preserva dois de cada quatro pesos, o que reduz pela metade o volume que o hardware precisa mover.
O que separa o MaskLLM é a forma como ele escolhe quais pesos manter. Métodos antigos se apoiam numa regra fixa, como descartar os menores pesos ou usar uma estimativa estatística. O MaskLLM trata a escolha da máscara como algo que o modelo aprende. Ele modela os padrões N:M como uma distribuição de probabilidade e a treina com uma técnica de amostragem chamada Gumbel Softmax, o que permite ao modelo encontrar máscaras melhores em um conjunto de dados amplo. Os pesos ficam congelados enquanto as máscaras são aprendidas, e é isso que torna a esparsidade 2:4 "sem perdas" para uma determinada tarefa.
O apelo prático está nos números. Segundo a NVIDIA, o modelo denso LLaMA-2 7B marca perplexidade de 5,12 no Wikitext, enquanto métodos de poda anteriores empurram esse valor para além de 10. O MaskLLM chega a 6,72 ou 6,78 apenas aprendendo as máscaras, com os pesos congelados. O projeto é de código aberto sob uma licença da NVIDIA, então pesquisadores e engenheiros podem rodá-lo por conta própria. Ainda assim, exige hardware com GPU de verdade e conhecimento de aprendizado de máquina.
Getting Started
- Instale os pacotes de Python e rode o script que baixa uma máscara já calculada do Hugging Face
- Carregue o modelo desejado, como o LLaMA-2 7B ou o LLaMA-3 8B, no PyTorch
- Aplique a máscara baixada e meça a precisão com o script de perplexidade incluído no projeto
- Para trabalho próprio, monte o ambiente do Megatron-LM, prepare os dados e treine uma máscara nova no seu domínio
Informações do produto
Uma visão rápida dos preços, das plataformas compatíveis e do desempenho de MaskLLM.
Ideal para
Os usuários, tarefas e cenários em que esta ferramenta se encaixa melhor.
Usuários
- Pesquisadores de aprendizado de máquina
- Engenheiros de ML que fazem deploy de modelos
- Estudantes e acadêmicos
- Desenvolvedores de frameworks e ferramentas
Tarefas
- Poda de modelos de linguagem
- Compressão de um único passo
- Esparsidade para um domínio específico
- Transferência de esparsidade
- Comparação de métodos de poda
Cenários
- Uma equipe que serve um modelo de 7B com poucas GPUs
- Pesquisa que precisa reproduzir números publicados
- Um modelo bom no geral que falha num nicho
- Montar um pipeline de compressão para uma família de modelos
- Estudar limites de licença ou privacidade
Principais recursos
Máscaras N:M aprendíveis
Em vez de pontuar cada peso à mão, o MaskLLM aprende a máscara. Ele trata o conjunto de padrões N:M válidos como uma distribuição e faz amostragens durante o treino, então o modelo converge para as melhores máscaras de um conjunto de dados. Esse passo aprendível é a diferença de fundo em relação aos métodos de um único passo, que escolhem uma máscara uma vez e param.
Amostragem com Gumbel Softmax
O método toma emprestada uma técnica de amostragem discreta para tornar a escolha da máscara treinável de ponta a ponta. O Gumbel Softmax deixa os gradientes passarem por aquilo que, de outro modo, seria uma escolha dura, de tudo ou nada. É isso que permite otimizar a distribuição sobre as máscaras N:M com dados reais em vez de estimá-la.
Máscaras já calculadas para modelos do Hugging Face
O projeto já vem com máscaras prontas para você testar sem treinar nada. Há máscaras publicadas para o LLaMA-2 7B, o LLaMA-3 8B e o LLaMA-3.1 8B, cada uma treinada num subconjunto do C4, e um script baixa todas do Hugging Face e as aplica aos modelos oficiais para avaliação.
Esparsidade semiestruturada 2:4
O MaskLLM mira a esparsidade N:M, quase sempre no padrão 2:4 que preserva dois valores diferentes de zero em cada quatro pesos consecutivos. Esse formato importa porque casa com a aceleração por hardware, então o modelo menor não ocupa só menos espaço: também pode rodar mais rápido na GPU certa.
Transferência de padrões esparsos
O mesmo modelo de probabilidade que aprende uma máscara consegue levá-la de um domínio para outro. Depois de treinar uma distribuição de máscaras, ela se adapta a uma tarefa ou conjunto de dados novo com menos esforço do que treinar do zero, o que os autores chamam de transferibilidade. Isso ajuda quando um modelo precisa continuar afiado em mais de um cenário.
Pipeline de treinamento com Megatron-LM
Para treinamentos completos, o MaskLLM inclui um fluxo com o Megatron-LM que divide um modelo LLaMA entre várias GPUs. O repositório traz scripts para converter checkpoints, tokenizar dados, gerar uma máscara prévia com SparseGPT, Wanda ou poda por magnitude, treinar a máscara e cortar o checkpoint antes de publicar.
Prós e contras
Prós
- Aprende as máscaras a partir dos dados em vez de depender de uma regra de importância fixa, o que melhora a precisão na mesma esparsidade
- Publica código, pesos e máscaras já calculadas sob licença aberta, então os resultados podem ser reproduzidos
- Aceita aprendizado por transferência, então uma distribuição de máscaras serve para várias tarefas e domínios
- Funciona com modelos conhecidos como LLaMA-2, LLaMA-3 e LLaMA-3.1 por meio do Hugging Face e do Megatron-LM
- Mantém os pesos congelados, então a esparsidade 2:4 aprendida não tem perdas para a tarefa em que foi ajustada
Contras
- Treinar uma máscara nova exige um equipamento com várias GPUs, e o repositório aponta cerca de 40 GB de memória por GPU nas execuções completas
- É um artefato de pesquisa, não um produto: não há suporte, nem serviço hospedado, nem API
- As máscaras prontas saem para poucos modelos LLaMA, então outras arquiteturas obrigam você a treinar a sua
- O ganho aparece em hardware que aceita a aceleração 2:4, o que limita o que GPUs antigas ou não NVIDIA aproveitam
Perguntas frequentes
O MaskLLM é um método de poda aprendível que adiciona esparsidade N:M a modelos de linguagem grandes. Em vez de pontuar cada peso, ele aprende quais manter modelando as máscaras como uma distribuição de probabilidade. A NVIDIA Research o publicou na NeurIPS 2024 e o código do MaskLLM é aberto no GitHub.
Conteúdo relacionado
Explore ferramentas, skills e artigos relacionados a MaskLLM.
Alternativas a MaskLLM
Base Code
Base44 (Wix) · Programação · ProdutividadeO Base Code é um ambiente de desenvolvimento em nuvem compartilhado da Base44. O time descreve mudanças em um repositório do GitHub pelo chat, confere tudo na prévia e envia o resultado como pull request.
Google EmbeddingGemma 2
Google DeepMind · Escrita · Imagem · ProgramaçãoGoogle EmbeddingGemma 2 é um modelo aberto de 740 milhões de parâmetros que transforma texto, código, imagens, áudio e quadros de vídeo em vetores de um único espaço de 768 dimensões, rodando em celulares e notebooks.
GitGlow
GitGlow · Programação · ProdutividadeO GitGlow é um cliente Git de desktop que coloca seu código e as mudanças dos seus agentes de IA na mesma revisão, para você saber o que vai sair da máquina antes de publicar.