Pruna-AI – Querem correr melhores modelos de IA em hardware mais modesto?

Higgs

Olá a todos.

Quem daqui tem uma placa gráfica com 8 GB de ram e quer correr modelos com mais complexidade e que requerem memoria vram até ao fim dos tempos levante a mão!
Todos nós temos este problema. Uns com mais intensidade, outros com menos, mas todos os entusiastas de LLM que conheço tem este problema: Falta de vram para correr os seus modelos.

A sentir este problema fortemente na pele, decidi procurar no huggingface modelos destilados ou optimizados para ajudar a sanar a falha. Notei que existiam vários, sendo que a maior parte deles vinham de uma empresa/group chamado Pruna. Experimentei vários e notei que embora fossem dumber que os seus modelos não destilados, se os usasse apenas para a tarefa para a qual foram desenhados, eram surpreendente mais leves em recursos e apresentavam resultados muito aceitáveis.
Decidi então procurar o que é isto do Pruna-AI, e cheguei ao post desta semana: O Pruna-AI é uma ferramenta inovadora projetada para otimizar modelos de linguagem de grande escala (LLMs), reduzindo o seu tamanho sem comprometer significativamente o desempenho. Esta técnica, conhecida como “pruning” ou poda, é essencial para adaptar modelos complexos às limitações de hardware, como as placas gráficas e aceleradores de IA disponíveis.
Para este post iremos utilizar o Docker para fazer deployment do produto, onde o podemos instalar e executar de forma eficiente e isolada, garantindo um ambiente consistente e reproduzível.

O Pruna-AI é uma ferramenta de código aberto que implementa algoritmos avançados de poda para modelos de aprendizagem profunda. A poda é uma técnica que remove pesos ou neurónios menos significativos de uma rede neural, resultando em modelos mais leves e eficientes. Esta abordagem é particularmente útil para LLMs, que frequentemente possuem bilhões de parâmetros, tornando-os difíceis de implementar em hardware com recursos limitados.

Pruna AI · GitHub

Mas existem benefícios da Poda de Modelos LLM?

Sim existem, da mesma forma que existem downsides. Já dizia o Einstein, se avaliares um peixe pela capacidade de subir a uma arvore será sempre um idiota.
Os upsides e downsides que encontrei com as minhas experiencias foram:

Os benefícios da poda de modelos LLM incluem:

  • Redução de requisitos computacionais – Modelos podados requerem menos memória e recursos de processamento, permitindo a sua execução em dispositivos com recursos limitados.
  • Maior velocidade de inferência – Com menos parâmetros para processar, os modelos podados geralmente oferecem tempos de resposta mais rápidos.
  • Menor consumo energético – Modelos mais pequenos consomem significativamente menos energia, reduzindo custos operacionais e impacto ambiental.
  • Custos de implementação reduzidos – Hardware menos potente e infraestrutura mais simples são necessários para modelos podados, reduzindo custos de implementação.
  • Democratização da tecnologia – Permite que organizações com recursos limitados utilizem modelos LLM que de outra forma seriam inacessíveis.
  • Potencial para melhor generalização – Em alguns casos, a remoção de parâmetros redundantes pode ajudar a reduzir o sobreajuste (overfitting) e melhorar a generalização.
  • Facilita implementações em dispositivos edge – Possibilita a utilização em telemóveis, dispositivos IoT e outros sistemas com restrições de recursos.
  • Latência reduzida – Especialmente importante para aplicações em tempo real onde o tempo de resposta é crítico.
  • Possibilidade de especialização – A poda pode ser direcionada para preservar capacidades específicas, criando modelos mais eficientes para tarefas específicas.
  • Menor pegada de armazenamento – Os modelos podados ocupam menos espaço em disco e memória, simplificando a distribuição e armazenamento.

Os downsides que encontrei com as minhas experiencias foram:

  • Modelos mais Específicos Podem ser Piores: Modelos menores ocupam menos espaço de armazenamento, no entanto estamos a desbastar neurónios de uma rede neural, basicamente dumb’ing it down.
  • Degradação do desempenho – A poda remove inerentemente parâmetros, o que pode levar a uma redução da qualidade do modelo em várias tarefas. Isto manifesta-se tipicamente como menor precisão, capacidades de raciocínio reduzidas ou qualidade de geração degradada.
  • Impacto desproporcional em capacidades específicas – A poda afeta frequentemente algumas capacidades do modelo mais do que outras. Por exemplo, um modelo podado pode manter fortes capacidades de resumo, mas mostrar uma degradação substancial no raciocínio complexo ou no conhecimento de domínios especializados.
  • Potencial amplificação de enviesamentos – Se não for implementada cuidadosamente, a poda pode inadvertidamente amplificar enviesamentos existentes no modelo ao remover parâmetros que ajudam a equilibrar certas representações.
  • Recuperação limitada de conhecimento – Uma vez removidos os parâmetros, essa codificação específica de conhecimento é perdida e não pode ser facilmente recuperada sem retreino completo.
  • Requisitos complexos de avaliação – Determinar quais parâmetros podar requer avaliação sofisticada em diversas tarefas para garantir que capacidades críticas não sejam perdidas.
  • Retornos decrescentes – Existe tipicamente uma relação não linear entre a taxa de poda e a perda de desempenho. Para além de um certo limiar (frequentemente 30-50%, dependendo do modelo), o desempenho cai precipitadamente.
  • Constrangimentos de arquitetura – Alguns métodos de poda só funcionam bem com arquiteturas específicas de modelos ou requerem suporte de hardware/framework para concretizar os seus ganhos teóricos de eficiência.
  • Sobrecarga de retreino – Muitas abordagens eficazes de poda requerem um ajuste fino pós-poda para recuperar o desempenho, adicionando custo computacional ao processo.

Instalação do Pruna-AI com Docker

Agora que temos uma noção do bem ou mail que isto pode correr, vamos meter a mão na massa. Para instalar e utilizar o Pruna-AI utilizando o Docker, efetuem os passos abaixo:

Nota: Se ainda não tiverem o Docker instalado, podem obtê-lo a partir do site oficial: https://www.docker.com/get-started. Sigam as instruções específicas para o vosso sistema operativo.

Criar o Dockerfile

Primeiro, criem um Dockerfile com as dependências necessárias:

FROM pytorch/pytorch:2.0.0-cuda11.7-cudnn8-runtime

# Instalar dependências básicas
RUN apt-get update && apt-get install -y \
    git \
    wget \
    curl \
    build-essential \
    && rm -rf /var/lib/apt/lists/*

# Configurar ambiente de trabalho
WORKDIR /app

# Copiar requisitos e instalar dependências Python
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Instalar Pruna-AI
RUN pip install pruna-ai
# Alternativa: instalação do código fonte
# RUN git clone https://github.com/VectorInstitute/pruna.git && \
#     cd pruna && \
#     pip install -e .

# Copiar scripts de pruning para o container
COPY scripts/ /app/scripts/
COPY data/ /app/data/

# Definir um volume para persistir modelos podados
VOLUME /app/models

# Comando padrão para executar o script de pruning
ENTRYPOINT ["python", "/app/scripts/run_pruning.py"]

Criar o ficheiro requirements.txt

transformers>=4.30.0
datasets>=2.12.0
accelerate>=0.20.0
torch>=2.0.0
peft>=0.4.0
evaluate>=0.4.0
tensorboard>=2.13.0
scipy>=1.10.0
scikit-learn>=1.2.2

Criar o Script de Pruning

Criem  um ficheiro scripts/run_pruning.py:

import os
import argparse
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
from datasets import load_dataset
from pruna import PruningScheduler, PruningConfig
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def parse_args():
    parser = argparse.ArgumentParser(description="Script para pruning de modelos LLM com Pruna-AI")
    parser.add_argument("--model_name", type=str, required=True, help="Nome ou caminho do modelo base")
    parser.add_argument("--output_dir", type=str, default="/app/models/pruned_model", help="Diretório para salvar o modelo podado")
    parser.add_argument("--dataset_name", type=str, default="wikitext", help="Nome do dataset para fine-tuning/avaliação")
    parser.add_argument("--dataset_config", type=str, default="wikitext-2-raw-v1", help="Configuração do dataset")
    parser.add_argument("--pruning_method", type=str, default="magnitude", help="Método de pruning (magnitude, movement, taylor)")
    parser.add_argument("--target_sparsity", type=float, default=0.3, help="Taxa alvo de sparsity (0.0-1.0)")
    parser.add_argument("--pruning_epochs", type=int, default=5, help="Número de épocas para aplicar pruning gradualmente")
    parser.add_argument("--batch_size", type=int, default=8, help="Tamanho do batch para treino")
    parser.add_argument("--learning_rate", type=float, default=5e-5, help="Taxa de aprendizagem")
    parser.add_argument("--num_epochs", type=int, default=10, help="Número total de épocas de treino")
    return parser.parse_args()

def main():
    args = parse_args()
    logger.info(f"Iniciando pruning para o modelo {args.model_name}")
    
    # Verificar disponibilidade de GPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logger.info(f"Usando dispositivo: {device}")
    
    # Carregar modelo e tokenizer
    logger.info("Carregando modelo e tokenizer...")
    model = AutoModelForCausalLM.from_pretrained(args.model_name)
    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
    
    # Configurar tokenizer
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    # Carregar dataset
    logger.info(f"Carregando dataset {args.dataset_name}/{args.dataset_config}...")
    dataset = load_dataset(args.dataset_name, args.dataset_config)
    
    # Pré-processar dataset
    def tokenize_function(examples):
        return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)
    
    tokenized_dataset = dataset.map(tokenize_function, batched=True)
    
    # Dividir dataset
    if "train" in tokenized_dataset:
        train_dataset = tokenized_dataset["train"]
        eval_dataset = tokenized_dataset["validation"] if "validation" in tokenized_dataset else tokenized_dataset["test"]
    else:
        # Se não houver divisão predefinida, criar uma
        splits = tokenized_dataset["train"].train_test_split(test_size=0.1)
        train_dataset = splits["train"]
        eval_dataset = splits["test"]
    
    # Configurar pruning
    logger.info(f"Configurando pruning com método {args.pruning_method} e sparsity alvo {args.target_sparsity}...")
    pruning_config = PruningConfig(
        method=args.pruning_method,
        target_sparsity=args.target_sparsity,
        scheduler="linear",
        pruning_epochs=args.pruning_epochs
    )
    
    pruning_scheduler = PruningScheduler(
        model=model,
        config=pruning_config
    )
    
    # Configurar treinamento
    training_args = TrainingArguments(
        output_dir=args.output_dir,
        num_train_epochs=args.num_epochs,
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=args.batch_size,
        learning_rate=args.learning_rate,
        weight_decay=0.01,
        logging_dir=f"{args.output_dir}/logs",
        logging_steps=50,
        evaluation_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
        push_to_hub=False,
    )
    
    # Personalizar Trainer para incluir pruning
    class PruningTrainer(Trainer):
        def training_step(self, model, inputs):
            loss = super().training_step(model, inputs)
            pruning_scheduler.step()
            return loss
    
    trainer = PruningTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        tokenizer=tokenizer,
    )
    
    # Executar treinamento com pruning
    logger.info("Iniciando treinamento com pruning...")
    trainer.train()
    
    # Avaliar modelo final
    logger.info("Avaliando modelo final...")
    eval_results = trainer.evaluate()
    logger.info(f"Resultados da avaliação: {eval_results}")
    
    # Obter estatísticas de sparsity
    sparsity_stats = pruning_scheduler.get_sparsity_stats()
    logger.info(f"Sparsity global final: {sparsity_stats['global_sparsity']:.2%}")
    
    # Salvar modelo podado
    logger.info(f"Salvando modelo podado em {args.output_dir}...")
    trainer.save_model(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)
    
    # Opcional: Exportar modelo em formato otimizado
    optimized_output_dir = f"{args.output_dir}_optimized"
    logger.info(f"Exportando modelo otimizado para {optimized_output_dir}...")
    pruning_scheduler.export_pruned_model(optimized_output_dir, format="sparse")
    
    logger.info("Processo de pruning concluído com sucesso!")

if __name__ == "__main__":
    main()

Criar Script de Execução Docker

Criem um script shell run_docker_pruning.sh:

#!/bin/bash

# Definir variáveis
MODEL_NAME="EleutherAI/gpt-neo-1.3B"
OUTPUT_DIR="./pruned_models"
DATASET_NAME="wikitext"
DATASET_CONFIG="wikitext-2-raw-v1"
PRUNING_METHOD="magnitude"
TARGET_SPARSITY=0.3
PRUNING_EPOCHS=5
NUM_EPOCHS=10
BATCH_SIZE=8
LEARNING_RATE=5e-5

# Garantir que o diretório de saída existe
mkdir -p $OUTPUT_DIR

# Construir a imagem Docker
docker build -t pruna-ai-pruning .

# Executar o container
docker run --gpus all \
  -v $(pwd)/pruned_models:/app/models \
  --name pruna_pruning \
  pruna-ai-pruning \
  --model_name $MODEL_NAME \
  --output_dir /app/models/pruned_model \
  --dataset_name $DATASET_NAME \
  --dataset_config $DATASET_CONFIG \
  --pruning_method $PRUNING_METHOD \
  --target_sparsity $TARGET_SPARSITY \
  --pruning_epochs $PRUNING_EPOCHS \
  --num_epochs $NUM_EPOCHS \
  --batch_size $BATCH_SIZE \
  --learning_rate $LEARNING_RATE

Estrutura de Diretórios

Para a organização do nosso projeto temos:

pruna-docker/
├── Dockerfile
├── requirements.txt
├── run_docker_pruning.sh
├── scripts/
│   └── run_pruning.py
├── data/
│   └── [optional data files]
└── pruned_models/  # Será criado automaticamente

Executar o Processo de Pruning

# Dar permissão de execução ao script
chmod +x run_docker_pruning.sh

# Executar o script
./run_docker_pruning.sh

Monitorizar o Progresso

Para ver os logs do processo de pruning em execução:

docker logs -f pruna_pruning 

Para aceder ao TensorBoard e visualizar métricas (em outro terminal/sessão):

docker exec -it pruna_pruning tensorboard --logdir=/app/models/pruned_model/logs --host 0.0.0.0

Depois acedam a http://localhost:6006 no vosso browser.

Recuperar o Modelo Podado

O modelo podado estará disponível no diretório ./pruned_models/pruned_model na vossa máquina host, graças ao volume Docker que configurámos previamente.

Opções Avançadas:

E se tivermos ou quisermos distribuir carga por vários GPUs? Criem um arquivo docker-compose.yml:

version: '3'

services:
  pruning:
    build: .
    volumes:
      - ./pruned_models:/app/models
    command: >
      --model_name EleutherAI/gpt-neo-1.3B
      --output_dir /app/models/pruned_model
      --pruning_method magnitude
      --target_sparsity 0.3
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]

E executem-o com:

docker-compose up 

Avaliação do Modelo Podado

Para avaliar o modelo após o pruning, criem outro script evaluate_pruned_model.py:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
import evaluate
import argparse

def parse_args():
    parser = argparse.ArgumentParser(description="Avaliação de modelo LLM podado")
    parser.add_argument("--model_path", type=str, required=True, help="Caminho para o modelo podado")
    parser.add_argument("--dataset_name", type=str, default="wikitext", help="Nome do dataset de avaliação")
    parser.add_argument("--dataset_config", type=str, default="wikitext-2-raw-v1", help="Configuração do dataset")
    return parser.parse_args()

def main():
    args = parse_args()
    
    # Carregar modelo e tokenizer
    model = AutoModelForCausalLM.from_pretrained(args.model_path)
    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
    
    # Carregar dataset
    dataset = load_dataset(args.dataset_name, args.dataset_config, split="test")
    
    # Avaliar perplexidade
    perplexity = evaluate.load("perplexity")
    results = perplexity.compute(
        predictions=dataset["text"][:100],
        model_id=args.model_path,
    )
    
    print(f"Perplexidade: {results['mean_perplexity']}")
    
    # Adicione outras métricas conforme necessário

if __name__ == "__main__":
    main()

Execute a avaliação com:

docker run --gpus all \
  -v $(pwd)/pruned_models:/app/models \
  pruna-ai-pruning \
  python /app/scripts/evaluate_pruned_model.py \
  --model_path /app/models/pruned_model

Melhores Práticas e Considerações

Nesse caso, quais as melhores praticas e considerações antes de ir vender a ideia que isto funciona e pode ajudar a baixar o custo de corrermos LLM’s nas nossas empresas ou nos nossos laboratórios?

  • Backup dos Modelos Originais: Antes de realizar qualquer poda, por favor, façam cópias de segurança dos modelos originais para evitar perdas irreversíveis. Não deveria ser preciso dizer, mas digo á mesma.
  • Testes A/B: Comparem o desempenho dos modelos podados com os originais em cenários reais (e com a mesma pergunta) para avaliar a viabilidade da otimização.
  • Poda Progressiva: Em vez de remover uma grande percentagem de pesos de uma só vez, façam a poda em etapas menores e avalie os impactos em cada fase. Isto ajuda a manter a precisão do modelo.
  • Re-treinamento Pós-Poda: Em alguns casos, após a remoção de pesos, um breve re-treinamento pode ajudar a recuperar parte da precisão perdida.
  • Ajuste Fino dos Parâmetros: Diferentes modelos vão reagir de forma distinta à poda. Experimente variações na percentagem de sparsity para encontrar o equilíbrio ideal entre tamanho e desempenho.
  • Monitoramento de Recursos: Ao executar modelos podados, monitorizem o uso de GPU, memória e tempo de inferência para garantir que a otimização trouxe benefícios reais. A ideia é diminuição aceitável versus poupança. Não nos interessa diminuição não aceitável se ficarmos a gastar os mesmos recursos.
  • Exploração de Diferentes Técnicas de Poda: Além da poda baseada em sparsity, explore métodos como poda estrutural (removendo neurônios inteiros) ou quantização para obter ainda mais eficiência.

Com estas práticas, é possível maximizar os benefícios do Pruna-AI e garantir que a otimização dos modelos seja eficaz sem comprometer a precisão e a utilidade dos nossos modelos favoritos de IA.

E chegamos ao fim de mais um post semanal. Neste exploramos como o Pruna-AI é uma ferramenta poderosa para otimizar modelos de linguagem de grande escala (LLMs), permitindo reduzir significativamente o seu tamanho sem comprometer demasiado o desempenho. Ao utilizar técnicas de poda, conseguimos melhorar a eficiência da inferência, reduzir o consumo de memória e energia, e tornar modelos pesados mais viáveis para execução em hardware limitado.
A instalação e utilização do Pruna-AI com Docker facilita enormemente a configuração e garante um ambiente estável e isolado para a otimização de modelos. Com apenas alguns comandos, é possível analisar, podar e testar modelos, maximizando a utilização das placas gráficas e aceleradores de IA disponíveis e ao adotar estas práticas, não só conseguimos tornar os LLMs mais acessíveis e eficientes, como também reduzimos os custos operacionais e melhoramos a escalabilidade de aplicações baseadas em IA. Para quem trabalha com modelos de inteligência artificial, a poda é uma estratégia essencial para equilibrar desempenho e recursos computacionais. Win-Win.

Até ao post da próxima semana, e já sabem, se tiverem alguma sugestão, ou correção sabem onde me encontrar.
Um abraço
Nuno

Related Posts