Pitch: Como ganhei convergência de 2ª Ordem "de graça" no Fine-Tuning de LLMs (Apresentando o SCAO)
Fala, galera! Queria compartilhar com vcs um projeto open-source focado em infraestrutura de Machine Learning que acabei de publicar e que está em processo de integração em grandes frameworks como o Axolotl.
Se vc faz fine-tuning de LLMs hoje, quase com certeza vc usa o AdamW (ou suas variantes 8-bit/Paged). Ele é estável e eficiente. Mas ele tem um defeito fundamental: ele é cego para a curvatura do espaço de erro (é um otimizador de 1ª ordem).
Nós sabemos que otimizadores de 2ª ordem convergem muito mais rápido e com menos passos. O problema é que, no mundo real da engenharia, eles são um tabu para treinamento local. A matemática pesada das matrizes Hessianas explode a VRAM da sua placa de vídeo e derruba a velocidade (throughput) pro o chão.
Como desenvolvedor backend mergulhando em otimização de modelos, decidi atacar esse problema. Foi assim que criei o SCAO (Sparse Curvature-Aware Adaptive Optimizer).
O "Pulo do Gato": Explorando o Gargalo de Memória
A grande sacada do SCAO não é apenas uma matemática mais limpa, mas o uso inteligente de um gargalo físico do hardware.
Quando você treina um modelo de bilhões de parâmetros usando QLoRA (pesos em 4-bit, adaptadores em 16-bit/32-bit), o treinamento não é limitado pelo processamento (Compute-Bound). Ele é limitado pela memória (Memory-Bound).
Basicamente, a sua GPU tem núcleos de processamento (CUDA cores) super rápidos esperando os pesos mastigados chegarem por uma esteira lenta (VRAM bandwidth).
O SCAO se aproveita desse tempo ocioso. Enquanto a GPU está "esperando" os pesos do modelo base carregarem da VRAM, o SCAO calcula as pesadas matrizes de curvatura esparsa apenas nos parâmetros treináveis do LoRA, rodando no fundo (em background).
O resultado? A matemática de 2ª ordem roda praticamente "de graça".
O Benchmark no "Metal Nu"
Para provar que funciona no mundo real e não só em clusters de milhões de dólares, rodei um benchmark limpo num hardware acessível que muita gente usa: uma simples Nvidia T4 (16GB) no Google Colab.
O Setup:
- Modelo:
Qwen/Qwen2.5-3B - Técnica: 4-bit QLoRA (r=16, alpha=32)
- Passos: 100 primeiros steps
| Otimizador | Throughput (Velocidade) | Loss (Menor é melhor) |
|---|---|---|
| AdamW (Baseline) | 16.65 s/it | 2.741 |
| SCAO (O meu) | 16.52 s/it | 2.523 |
Não só o SCAO entregou uma convergência mais rápida e inteligente (Loss muito menor nos mesmos 100 passos), como rodou um fio de cabelo mais rápido que o AdamW. Sem explosão de gradiente e sem crashes por falta de VRAM (OOM).
Como usar nos seus projetos
O SCAO é escrito em Python puro, herda de torch.optim.Optimizer e não exige compilação de C++ ou Ninja para rodar.
Já subi o pacote oficial no PyPI:
pip install scao
E para usar no seu código PyTorch, é um drop-in replacement:
import torch
from scao import SCAO
# Troque o torch.optim.AdamW por:
optimizer = SCAO(
model.parameters(),
lr=2e-4,
betas=(0.9, 0.999),
weight_decay=0.01
)
Próximos Passos e Open-Source
O SCAO já está com um Pull Request aberto no repositório oficial do Axolotl para ser suportado nativamente nos arquivos .yml, e estou abrindo as mesmas propostas para o LLaMA-Factory e o Torchtune da Meta.
Se vc treina modelos localmente e sofre com tempo de convergência ou limite de hardware, adoraria que vc testasse o pacote, tentasse quebrar o código e me desse um feedback.
Bora construir ferramentas de IA melhores! Qualquer dúvida sobre a implementação ou sobre QLoRA, deixem aí nos comentários que eu respondo.