Ainda se desespera com a lentidão do processamento de textos longos por modelos grandes? Não se preocupe! A Universidade Tsinghua lançou uma tecnologia revolucionária — o framework de raciocínio paralelo de sequência APB, que adiciona um motor de "turbocompressor" aos modelos grandes! Testes mostraram que esta tecnologia de ponta processa textos extremamente longos até 10 vezes mais rápido que o Flash Attention! Sim, você leu certo, 10 vezes!
Com o sucesso estrondoso de modelos grandes como o ChatGPT, a capacidade de "leitura" da IA aumentou significativamente, lidando facilmente com textos de mais de 100.000 palavras. No entanto, diante de uma quantidade enorme de informações, o "cérebro" dos modelos grandes tradicionais trava — a arquitetura Transformer, embora poderosa, tem um mecanismo de atenção central que funciona como um "super scanner". Quanto maior o texto, maior a expansão exponencial da área de escaneamento, e consequentemente, mais lenta a velocidade.
Para resolver este problema crítico, cientistas da Universidade Tsinghua, em colaboração com várias instituições de pesquisa e gigantes da tecnologia, encontraram uma solução inovadora: o framework APB. O segredo deste framework reside na combinação inteligente de "paralelismo de sequência + atenção esparsa".
Simplificando, o framework APB funciona como uma equipe de "trabalho colaborativo" eficiente. Ele "divide" o texto longo em pequenos pedaços, distribuindo-os entre vários GPUs que trabalham em paralelo. Ainda mais impressionante, o APB fornece a cada GPU habilidades de "compressão de cache KV local" e "comunicação simplificada", permitindo que elas processem suas tarefas individualmente e compartilhem informações-chave de forma eficiente, colaborando para resolver problemas complexos de dependência semântica em textos longos.
Surpreendentemente, o framework APB não sacrifica o desempenho em troca de velocidade. Pelo contrário, em testes com textos extremamente longos de 128K, o APB não apenas aumentou a velocidade, mas também superou o Flash Attention tradicional em termos de desempenho! Até mesmo o Star Attention, promovido pela Nvidia, foi superado pelo APB, com uma melhoria de velocidade de 1,6 vezes, tornando-o um verdadeiro "Ás".
A aplicação mais direta desta tecnologia inovadora é a redução significativa do tempo de resposta do primeiro token no processamento de textos longos por modelos grandes. Isso significa que, no futuro, modelos grandes equipados com o framework APB poderão compreender instantaneamente e responder rapidamente a comandos longos e detalhados dos usuários, eliminando a longa espera de "carregando...".
Então, como o framework APB consegue essa incrível aceleração?
O framework APB entende profundamente o "ponto crítico" do processamento de textos longos — a carga computacional. A carga computacional do mecanismo de atenção tradicional é proporcional ao quadrado do comprimento do texto; textos longos são um "buraco negro" computacional. Para superar este gargalo, o framework APB utiliza duas "técnicas secretas":
Primeira técnica: Aumentar o paralelismo, para que "muitas mãos fazem a obra leve"
O framework APB aproveita ao máximo as vantagens do cálculo distribuído, distribuindo as tarefas de cálculo entre várias GPUs, como em uma "colaboração em equipe", aumentando naturalmente a eficiência. Especialmente no paralelismo de sequência, o framework APB mostra uma escalabilidade extremamente forte, sem restrições à estrutura do modelo, permitindo lidar facilmente com textos de qualquer tamanho.
Segunda técnica: Reduzir cálculos inúteis, para que "o melhor seja usado para o melhor"
O framework APB introduz um mecanismo de atenção esparsa, que não "faz tudo de uma vez", mas sim calcula a atenção "seletivamente". É como um especialista com "olhos de águia", que se concentra apenas nas informações-chave do texto, ignorando as partes irrelevantes, reduzindo assim significativamente a carga computacional.
No entanto, estas duas técnicas, "paralelismo" e "esparsidade", parecem simples, mas na verdade "escondem segredos". Como implementar cálculos de atenção esparsa eficientes em um framework de paralelismo de sequência? Esta é a verdadeira "base sólida" do framework APB.
Em um ambiente de paralelismo de sequência, cada GPU possui apenas parte das informações do texto. Implementar uma atenção esparsa com "percepção global" é como "um cego tocando um elefante", uma tarefa extremamente difícil. Métodos anteriores, como Star Attention e APE, ou sacrificavam o desempenho ou tinham aplicações limitadas, e não conseguiam resolver perfeitamente este problema.
O framework APB, por sua vez, evita habilmente a "armadilha da comunicação em larga escala", encontrando uma solução alternativa, construindo um mecanismo de atenção esparsa de baixa comunicação voltado para cenários de paralelismo de sequência. Os principais componentes deste mecanismo incluem:
Blocos de âncora (Anchor block) menores: Os blocos de âncora funcionam como um "sistema de navegação", guiando o mecanismo de atenção para as informações-chave. O framework APB inova reduzindo o tamanho dos blocos de âncora, tornando-os mais leves e flexíveis, reduzindo a sobrecarga computacional.
Blocos de passagem (Passing block) exclusivos: Os blocos de passagem são os componentes "essenciais" do framework APB. Eles resolvem habilmente os problemas de dependência semântica de longo alcance. Ao "compactar e empacotar" as informações-chave processadas pela GPU anterior e transmiti-las para a GPU subsequente, cada "membro da equipe" pode ter uma "visão geral", compreendendo o contexto do texto longo.
Compressão de contexto sensível à consulta: O framework APB também introduz um mecanismo "sensível à consulta", permitindo que o compressor de contexto "entenda a pergunta", filtrando e retendo informações-chave relevantes para a consulta com maior precisão, melhorando ainda mais a eficiência e a precisão.
Com base nestas "técnicas exclusivas", o framework APB constrói um fluxo de raciocínio fluido:
Segmentação de contexto: Distribui uniformemente o texto longo entre as várias GPUs e adiciona um bloco de âncora no início, "inserindo" a pergunta da consulta.
Compressão de contexto: Utiliza o Locret para introduzir cabeças de retenção e realizar "compressão inteligente" do cache KV.
Comunicação eficiente: Utiliza o operador AllGather para "transmitir" o cache KV compactado para as GPUs subsequentes, construindo o bloco de passagem.
Cálculo ultrarrápido: Utiliza o Flash Attention Kernel personalizado, juntamente com uma máscara de atenção otimizada, para realizar cálculos eficientes. O bloco de passagem "cumpre sua missão" após o cálculo e não participa de cálculos subsequentes.
Os resultados experimentais demonstram inequivocamente o desempenho excepcional do framework APB. Em testes com vários modelos, como Llama-3.1-8B-instruct, Qwen-2.5-14B-instruct e Yi-34B-200K, e vários benchmarks, como InfiniteBench e RULER, o framework APB superou todos os outros, alcançando o melhor equilíbrio entre desempenho e velocidade.
Vale destacar que, com o aumento do comprimento do texto, a vantagem de velocidade do framework APB torna-se ainda mais evidente, alcançando um efeito surpreendente de "quanto maior, mais rápido". O segredo por trás disso é que a carga computacional do framework APB é muito menor do que outros métodos, e a diferença aumenta com o aumento do comprimento do texto.
Uma análise mais aprofundada do tempo de preenchimento prévio mostra que a tecnologia de paralelismo de sequência por si só pode reduzir significativamente o tempo de cálculo da atenção e da FFN (rede neural feedforward). O mecanismo de atenção esparsa do framework APB comprime ainda mais o tempo de cálculo da atenção ao extremo. Em comparação com o Star Attention, o framework APB utiliza habilmente o bloco de passagem para transmitir dependências semânticas de longo alcance, reduzindo significativamente o tamanho do bloco de âncora e reduzindo efetivamente a sobrecarga adicional da FFN, alcançando o efeito perfeito de "ter o melhor dos dois mundos".
Ainda mais animador é que o framework APB demonstra uma compatibilidade excepcional, podendo se adaptar de forma flexível a diferentes ambientes distribuídos e escalas de modelos, mantendo um alto desempenho e alta eficiência em várias condições "severas".
É previsível que, com o surgimento do framework APB, o "gargalo" do raciocínio de textos longos por modelos grandes será quebrado, e o espaço de imaginação das aplicações de IA será expandido infinitamente. No futuro, seja em atendimento ao cliente inteligente, análise financeira, pesquisa científica ou criação de conteúdo, entraremos em uma nova era de IA "mais rápida, mais poderosa e mais inteligente"!
Endereço do projeto: https://github.com/thunlp/APB
Endereço do artigo: https://arxiv.org/pdf/2502.12085
