Nova técnica de destilação de conhecimento reduz custos de treino para LLMs

A destilação de conhecimento, um processo utilizado para comprimir modelos de linguagem de grande escala (LLMs) em versões mais compactas, recebeu uma otimização técnica significativa. Numa publicação recente no Hugging Face, investigadores da Multiverse Computing detalharam uma abordagem que visa tornar o treino de modelos estudantes mais acessível, permitindo que o processo seja executado com exigências de memória VRAM reduzidas, algo anteriormente considerado um desafio devido às necessidades computacionais extremas.

Tradicionalmente, a destilação de conhecimento exige que o modelo professor e o modelo estudante estejam carregados simultaneamente na memória durante o treino. Este método, conhecido como destilação online, obriga a que o professor realize uma passagem completa para gerar distribuições de probabilidade sobre todo o vocabulário em cada passo. Com modelos modernos que possuem biliões de parâmetros e vocabulários vastos, o consumo de memória atinge valores elevados, frequentemente ultrapassando os 250GB de VRAM, o que excede a capacidade de memória de GPUs individuais de topo.

A solução proposta baseia-se em duas alterações sistémicas. A primeira consiste na destilação offline, onde as saídas do professor são calculadas apenas uma vez e armazenadas em cache, focando-se nos 100 tokens mais prováveis por posição. Desta forma, o modelo professor não necessita de permanecer na memória durante o treino do estudante, permitindo a reutilização dos dados em múltiplas experiências. Esta mudança elimina a necessidade de recomputar o comportamento do professor a cada iteração, reduzindo a carga computacional.

A segunda inovação introduzida é uma função de perda de divergência Kullback-Leibler (KL) fundida e fragmentada. Ao contrário das implementações padrão, que constroem uma matriz densa de vocabulário por comprimento de sequência, este método processa a informação em blocos. Esta técnica evita a materialização de matrizes de grandes dimensões, mantendo o consumo de memória sob controlo mesmo em sequências longas. Segundo os autores, esta abordagem permite que o pico de utilização de memória seja reduzido para cerca de 128GB.

Estes avanços técnicos representam um contributo para a eficiência no desenvolvimento de modelos de linguagem. Ao reduzir a barreira de hardware, a técnica facilita a experimentação e a criação de modelos especializados. A capacidade de realizar a destilação de modelos complexos, incluindo aqueles com janelas de contexto extensas, sem a necessidade de infraestruturas de clusters de centenas de GPUs, poderá otimizar o fluxo de trabalho no ecossistema de modelos de código aberto.