¿Sigues volviéndote loco por la lentitud de los grandes modelos de lenguaje al procesar textos largos? ¡No te preocupes! La Universidad Tsinghua presenta una tecnología revolucionaria: el marco de inferencia en paralelo de secuencias APB, ¡que añade un motor de "turbocompresor" a los grandes modelos de lenguaje! Las pruebas muestran que esta tecnología de vanguardia procesa textos extremadamente largos hasta 10 veces más rápido que Flash Attention. ¡Sí, has oído bien, 10 veces más rápido!
Con el auge de modelos como ChatGPT, la capacidad de "lectura" de la IA ha mejorado significativamente, pudiendo procesar fácilmente textos de más de cien mil palabras. Sin embargo, ante la información masiva, el "cerebro" de los grandes modelos de lenguaje tradicionales se atasca: la arquitectura Transformer, aunque potente, tiene un mecanismo de atención central que funciona como un "escáner súper potente". Cuanto más largo es el texto, más se expande el rango de escaneo de forma exponencial, lo que reduce la velocidad.
Para resolver este problema, científicos de la Universidad Tsinghua, junto con varias instituciones de investigación y gigantes tecnológicos, han creado el marco APB. El secreto de este marco radica en la ingeniosa combinación de "paralelismo de secuencias y atención dispersa".
En pocas palabras, el marco APB es como un equipo de "trabajo en equipo" altamente eficiente. "Desmiembra" el texto largo en pequeños fragmentos, los asigna a varios "miembros del equipo" de GPU para su procesamiento paralelo. Lo que es aún más ingenioso es que APB proporciona a cada "miembro del equipo" las habilidades de "compresión de caché KV local" y "comunicación simplificada", permitiéndoles compartir información clave de manera eficiente mientras procesan sus propias tareas, y colaborando para resolver problemas complejos de dependencia semántica en textos largos.
Sorprendentemente, el marco APB no sacrifica el rendimiento por velocidad. De hecho, en las pruebas con textos extremadamente largos de 128K, APB no solo aumenta la velocidad, sino que también supera el rendimiento de Flash Attention tradicional. Incluso supera a Star Attention, impulsado por NVIDIA, con una mejora de velocidad del 160%, convirtiéndose en un verdadero "as".
Esta tecnología innovadora se aplica directamente para reducir significativamente el tiempo de respuesta del primer token en el procesamiento de textos largos por parte de los grandes modelos de lenguaje. Esto significa que, en el futuro, los grandes modelos de lenguaje que incorporen el marco APB podrán comprender instantáneamente y responder rápidamente a las largas instrucciones de los usuarios, despidiéndose de la larga espera de "cargando...".
Entonces, ¿cómo logra el marco APB una aceleración tan "increíble"?
El marco APB comprende el "punto débil" del procesamiento de textos largos: la cantidad de cálculos. La cantidad de cálculos del mecanismo de atención tradicional es proporcional al cuadrado de la longitud del texto; los textos largos son un "agujero negro" para los cálculos. Para superar este obstáculo, el marco APB utiliza dos "movimientos maestros":
Primer movimiento: aumentar el paralelismo, "muchas manos hacen el trabajo ligero"
El marco APB aprovecha al máximo las ventajas del cálculo distribuido, distribuyendo las tareas de cálculo en múltiples GPU, como si fuera un "trabajo en equipo", lo que aumenta la eficiencia. Especialmente en el paralelismo de secuencias, el marco APB muestra una fuerte escalabilidad, sin estar limitado por la estructura del modelo, pudiendo manejar textos de cualquier longitud.
Segundo movimiento: reducir los cálculos ineficaces, "usar bien los recursos"
El marco APB introduce un mecanismo de atención dispersa, no calcula la atención de "forma indiscriminada", sino de forma "selectiva". Es como un experto con "ojos de águila", que solo se centra en la información clave del texto e ignora las partes irrelevantes, reduciendo así significativamente la cantidad de cálculos.
Sin embargo, estos dos movimientos, "paralelismo" y "dispersión", parecen simples, pero en realidad "ocultan secretos". ¿Cómo se puede lograr un cálculo de atención dispersa eficiente en un marco de paralelismo de secuencias? Esta es la verdadera "dureza" del marco APB.
En un entorno de paralelismo de secuencias, cada GPU solo tiene acceso a parte de la información del texto. Lograr una atención dispersa con "percepción global" es como "el ciego y el elefante", una tarea muy difícil. Los métodos anteriores, como Star Attention y APE, o sacrificaban el rendimiento o tenían aplicaciones limitadas, sin resolver completamente este problema.
El marco APB evita ingeniosamente el "problema de la comunicación a gran escala", buscando una nueva ruta y construyendo un mecanismo de atención dispersa de baja comunicación orientado a escenarios de paralelismo de secuencias. Los componentes principales de este mecanismo incluyen:
Bloques de anclaje (Anchor block) más pequeños: los bloques de anclaje son como un "navegador" que guía el mecanismo de atención hacia la información clave. El marco APB reduce innovadoramente el tamaño de los bloques de anclaje, haciéndolos más ligeros y flexibles, y reduciendo los costos de cálculo.
Bloques de paso (Passing block) exclusivos: los bloques de paso son los componentes "clave" del marco APB, que resuelven ingeniosamente el problema de la dependencia semántica a larga distancia. Al "comprimir y empaquetar" la información clave procesada por la GPU anterior y transmitirla a la GPU posterior, cada "miembro del equipo" puede tener una "visión general", comprendiendo el contexto del texto largo.
Compresión de contexto sensible a la consulta: el marco APB también introduce un mecanismo "sensible a la consulta", que permite al compresor de contexto "comprender la pregunta" y seleccionar y conservar con mayor precisión la información clave relacionada con la consulta, mejorando aún más la eficiencia y la precisión.
Con estas "habilidades especiales", el marco APB construye un flujo de inferencia fluido:
Segmentación de contexto: el texto largo se distribuye uniformemente entre las GPU, y se agrega un bloque de anclaje al principio para "insertar" la pregunta de consulta.
Compresión de contexto: se utiliza Locret para introducir encabezados reservados y realizar una "compresión inteligente" de la caché KV.
Comunicación eficiente: mediante el operador AllGather, la caché KV comprimida se "transmite" a las GPU posteriores para construir el bloque de paso.
Cálculo de alta velocidad: se utiliza un núcleo Flash Attention especial, junto con una máscara de atención optimizada, para realizar cálculos eficientes. El bloque de paso se "retira" después de completar el cálculo y no participa en los cálculos posteriores.
Los resultados experimentales demuestran el excelente rendimiento del marco APB. En las pruebas realizadas en varios modelos como Llama-3.1-8B-instruct, Qwen-2.5-14B-instruct e Yi-34B-200K, y en varios puntos de referencia como InfiniteBench y RULER, el marco APB ha superado a todos los demás, logrando el mejor equilibrio entre rendimiento y velocidad.
Cabe destacar que, a medida que aumenta la longitud del texto, la ventaja de velocidad del marco APB se hace más evidente, logrando un efecto "cuanto más largo, más rápido". El secreto detrás de esto es que la cantidad de cálculos del marco APB es mucho menor que la de otros métodos, y la diferencia aumenta a medida que aumenta la longitud del texto.
Un análisis más profundo del tiempo de prellenado muestra que la tecnología de paralelismo de secuencias puede reducir significativamente el tiempo de cálculo de la atención y la FFN (red neuronal de avance). El mecanismo de atención dispersa del marco APB comprime aún más el tiempo de cálculo de la atención. En comparación con Star Attention, el marco APB utiliza ingeniosamente el bloque de paso para transmitir dependencias semánticas a larga distancia, reduciendo significativamente el tamaño del bloque de anclaje y reduciendo eficazmente los costos adicionales de la FFN, logrando un efecto "perfecto" donde se obtiene lo mejor de ambos mundos.
Lo más emocionante es que el marco APB muestra una excelente compatibilidad, pudiendo adaptarse flexiblemente a diferentes entornos distribuidos y escalas de modelos, manteniendo un alto rendimiento y eficiencia en diversas condiciones "rigurosas".
Es previsible que, con la aparición del marco APB, se eliminará el "cuello de botella" de la inferencia de textos largos de los grandes modelos de lenguaje, y el espacio de imaginación de las aplicaciones de IA se expandirá infinitamente. En el futuro, ya sea en el servicio de atención al cliente inteligente, el análisis financiero, la investigación científica o la creación de contenido, entraremos en una nueva era de la IA "más rápida, más potente e inteligente".
Dirección del proyecto: https://github.com/thunlp/APB
Dirección del artículo: https://arxiv.org/pdf/2502.12085