ParaBlock: Communication-Computation Parallel Block Coordinate Federated Learning for Large Language Models

Resumen Técnico: ParaBlock: Aprendizaje Federado de Bloque de Coordenadas con Computación-Comunicación en Paralelo para Modelos de Lenguaje de Gran Escala

Planteamiento del Problema
El Aprendizaje Federado (FL, por sus siglas en inglés) ha emergido como un paradigma de preservación de la privacidad para entrenar modelos a través de datos descentralizados. Recientemente, los esquemas de descenso de coordenadas de bloques federados (BCD) han ganado tracción para el entrenamiento de modelos a gran escala, permitiendo que los clientes actualicen solo un subconjunto (bloque) de los parámetros del modelo localmente en lugar de todo el modelo. Sin embargo, en la era de los Modelos de Lenguaje de Gran Escala (LLMs), incluso un solo bloque puede contener decenas o cientos de millones de parámetros. Esta escala crea un cuello de botella crítico: la latencia de comunicación requerida para transmitir estos grandes bloques entre los clientes y el servidor se vuelve significativa, particularmente para los clientes periféricos con recursos limitados.

Los métodos tradicionales de BCD federado operan bajo un diseño de un solo hilo (comunicación-computación-comunicación), donde los clientes deben esperar a que la transmisión del modelo se complete antes de iniciar el siguiente paso de computación. Si bien este tiempo de comunicación era insignificante para modelos más pequeños, ahora obstaculiza severamente la escalabilidad y la eficiencia del ajuste fino (fine-tuning) federado de LLMs, causando un tiempo de inactividad sustancial para los clientes.

Metodología: ParaBlock
Para abordar esta latencia, los autores proponen ParaBlock, un nuevo método de BCD federado que introduce una arquitectura de comunicación-computación en paralelo. A diferencia de los enfoques estándar de un solo hilo, ParaBlock establece dos hilos paralelos en los clientes locales:

1. Hilo de Computación: Realiza las actualizaciones de coordenadas de bloques locales (ajuste fino del bloque asignado) utilizando datos locales.
2. Hilo de Comunicación: Se sincroniza con el servidor para enviar las actualizaciones locales de la ronda anterior y recibir las actualizaciones globales agregadas.

Mecanismos Clave:

• Solapamiento (Overlapping): Los clientes ejecutan la computación local mientras se comunican simultáneamente con el servidor. Esto elimina el tiempo de inactividad típicamente causado por la espera de la transmisión de red.
• Obsolescencia y Corrección: Debido a que el hilo de comunicación tiene un retraso de una ronda (los clientes reciben las actualizaciones globales de la ronda $t-1$ mientras computan para la ronda $t$), los parámetros del modelo local pueden volverse inconsistentes con el modelo global. ParaBlock aborda esto mediante un mecanismo de corrección. Cuando un cliente actualiza su modelo local para la siguiente ronda, no solo aplica la nueva actualización local al bloque actual, sino que también corrige los parámetros del bloque anterior utilizando la actualización global recibida ($\Delta_{t-1}$) para reemplazar la actualización local obsoleta ($\Delta^i_{t-1}$).
• Consistencia Recursiva: El análisis teórico demuestra que esta corrección asegura que, tras una ronda global, los parámetros del modelo local para el bloque previamente actualizado se alineen exactamente con los parámetros del modelo global, mitigando efectivamente el impacto del retraso de una ronda.

Contribuciones Clave

1. Arquitectura Novedosa: La propuesta de ParaBlock, que reemplaza el diseño tradicional de FL de un solo hilo con un esquema paralelo de dos hilos diseñado específicamente para reducir la latencia de comunicación en el ajuste fino federado de LLMs.
2. Garantías Teóricas: Los autores proporcionan un análisis de convergencia riguroso para objetivos no convexos. Demuestran que ParaBlock logra una tasa de convergencia de $O(1/\sqrt{T})$, la cual es consistente con los métodos estándar de descenso de coordenadas de bloques federados. Esto confirma que las ganancias de eficiencia derivadas del paralelismo no comprometen el rendimiento de la convergencia teórica.
3. Validación Empírica: Se realizaron experimentos extensos para el ajuste fino de los modelos Llama 3-8B y Llama 3.2-3B en tareas de seguimiento de instrucciones generales (Alpaca-GPT4) y razonamiento matemático (MathInstruct/GSM8K).

Resultos

• Desempeño: ParaBlock mantiene un sólido desempeño en tareas descendentes (downstream), logrando puntuaciones comparables o superiores a las de las líneas base de BCD federado estándar (por ejemplo, FedBCD) y superando significativamente a varios métodos de Ajuste Fino Eficiente en Parámetros (PEFT) basados en FL (por ejemplo, FedIT, FLoRA) en tareas de seguimiento de instrucciones y razonamiento matemático.
• Eficiencia: El método reduce significativamente el tiempo de ejecución (wall-clock runtime). En escenarios con bajo ancho de banda de red (50M/s) y tamaños de lote pequeños, ParaBlock redujo el tiempo de ejecución en más de un 30% en comparación con las líneas base de coordinación de bloques como FedCyBGD. Comparado con los métodos basados en LoRA como FedIT, ParaBlock logró una reducción de tiempo de aproximadamente el 40% en la mayoría de las configuraciones.
• Robustez: Los estudios de ablación mostraron que el método sigue siendo efectivo bajo distribuciones de datos heterogéneas (non-IID) y en diferentes arquitecturas de modelos (incluyendo Qwen-2.5). El método también demostró compatibilidad con configuraciones de participación parcial en escenarios de cross-silo.

Significancia
El artículo afirma que ParaBlock ofrece una solución simple pero efectiva al cuello de botella de la comunicación inherente al BCD federado para LLMs. Al solapar la comunicación y la computación, permite el despliegue escalable del ajuste fino federado de modelos grandes en clientes con recursos limitados sin sacrificar las garantías de convergencia o la utilidad del modelo. El trabajo destaca que la reducción de la latencia de comunicación a través del paralelismo arquitectónico es una estrategia viable y necesaria para la próxima generación de entrenamiento de LLMs federados, distinta y ortogonal a las técnicas existentes como la cuantización o la poda (pruning).