Resumen Técnico: Fusión de Modelos para la Agregación en DiLoCo

Planteamiento del Problema

El entrenamiento de Bajo Comunicación Distribuida (DiLoCo) ha surgido como un método para entrenar grandes modelos de lenguaje a través de clústeres de cómputo mal conectados ("islas") minimizando la sincronización global. En DiLoCo, los trabajadores realizan múltiples pasos de optimización local (pasos internos) antes de agregar sus actualizaciones de parámetros (pseudo-gradientes) mediante un optimizador global. Aunque es efectivo para reducir los costos de comunicación, DiLoCo sufre una degradación de rendimiento a medida que aumenta el número de trabajadores ($R$) o el número de pasos locales ($H$). Esta degradación ocurre porque las réplicas locales divergen significamente, causando que el promedio de los pseudo-gradientes se desalinee con las actualizaciones producidas por el entrenamiento de paralelismo de datos estrictamente sincronizado.

Concurrentemente, se han desarrollado técnicas de Fusión de Modelos (Model Merging) para combinar modelos ajustados de forma independiente. Un desafío central en la fusión es la "interferencia de tareas", donde las actualizaciones de parámetros de diferentes tareas entran en conflicto, provocando caídas en el rendimiento, especialmente a medida que aumenta el número de modelos a fusionar.

Los autores identifican una analogía estructural entre estos dos campos: el paso de agregación de pseudo-gradientes en DiLoCo (con un optimizador externo SGD) es matemáticamente equivalente a la Aritmética de Tareas (Task Arithmetic) de la fusión de modelos. En esta analogía:

• El modelo global compartido al inicio de un paso externo corresponde a la inicialización pre-entrenada.
• Las réplicas locales de los trabajadores corresponden a modelos "expertos" ajustados.
• Las actualizaciones locales (pseudo-gradientes) corresponden a "vectores de tarea".
• El número de trabajadores corresponde al número de tareas.

El problema central abordado es si las técnicas avanzadas de fusión de modelos, diseñadas específicamente para mitigar la interferencia de tareas, pueden ser reutilizadas para mejorar el paso de agregación en DiLoCo, restaurando así el rendimiento a medida que aumentan el número de trabajadores o los pasos locales.

Metodología

Puente Teórico

El artículo establece que la actualización externa de DiLoCo con SGD es una aplicación directa de la aritmética de tareas. Al visualizar la agregación de pseudo-gradientes como un problema de fusión de modelos, los autores proponen sustituir la operación de promediado estándar por algoritmos de fusión sofisticados diseñados para reducir la interferencia negativa.

Evaluación Experimental de Métodos de Fusión

Los autores evaluaron sistemáticamente varios métodos de fusión de vanguardia como reemplazos directos para el paso de agregación de DiLoCo. Los métodos evaluados incluyeron:

• Basados en la dispersión (Sparsity-based): TIES (TrIm, Elect Sign) y DARE (Drop And REscale).
• Basados en la ortogonalización: TSV (Task Singular Vector), Iso-C (Isotropic merging in Common subspace) e Iso-CTS.

Hallazgos sobre los Métodos de Fusión:

• Métodos de dispersión (TIES, DARE): Se comportaron de manera similar al SGD estándar de DiLoCo solo en niveles de dispersión muy bajos. El rendimiento se degradó a medida que aumentó la dispersión, lo que sugiere que no ofrecieron una ventaja distintiva sobre el promediado simple en este contexto.
• TSV: Funcionó deficientemente (mayor pérdida de validación) y fue computacionalmente prohibitivo, requiriendo múltiples Descomposiciones en Valores Singulares (SVD) por paso.
• Iso-CTS: Superó al baseline pero fue computacionalmente costoso y no ofreció una ventaja significativa sobre Iso-C en el entorno de DiLoCo (donde las tareas locales son cortas y similares).
• Iso-C: Emergió como el candidato más prometedor. Mejora el rendimiento reemplazando el espectro de valores singulares del pseudo-gradiente agregado por su media, efectivamente aplanando las direcciones de alta energía para obtener un subespacio común equilibrado. Requiere solo una SVD por matriz de parámetros, ofreciendo una relación costo-rendimiento favorable.

Algoritmo Propuesto: IsoLoCo

Basándose en el éxito de Iso-C, los autores proponen IsoLoCo (Isotropic Distributed Low-Communication). IsoLoCo adapta Iso-C para el entrenamiento distribuido mediante:

1. Corrección Isotrópica: Para matrices de parámetros 2D (capas de peso), calcula la SVD del pseudo-gradiente promediado, reemplaza los valores singulares con su media y reconstruye la actualización. Esto regulariza el espectro mientras preserva la estructura del subespacio.
2. Momento de Nesterov: A diferencia del Iso-C estándar, IsoLoCo incorpora el momento de Nesterov en el optimizador externo. Esto está motivado por trabajos previos que muestran el beneficio del momento en entornos FedOpt/DiLoCo. El pseudo-gradiente corregido se introduce en un búfer de momento para generar la actualización final.

Optimización de la Implementación

Para abordar el costo computacional de la SVD, los autores desarrollaron una implementación rápida de IsoLoCo utilizando iteraciones de Newton–Schulz. Esto reemplaza la SVD estándar por una ortogonalización iterativa y utiliza un proxy RMS para la media de los valores singulares. Esta optimización produce una ortogonalización hasta dos órdenes de magnitud más rápida con una degradación de rendimiento insignificante.

Resultados Clave

Rendimiento a través de Conteos de Trabajadores

Los experimentos se realizaron en modelos tipo Llama (178M, 512M y 1B de parámetros) utilizando el conjunto de datos DCLM.

• Escalabilidad con Trabajadores: IsoLoCo supera significativamente al DiLoCo estándar a medida que aumenta el número de trabajadores.
  • En la escala de 178M con $R=128$, IsoLoCo redujo la brecha de pérdida de validación respecto al baseline de paralelismo de datos en un 5.8% comparado con DiLoCo.
  • Con $R=128$, IsoLoCo mostró una degradación del 9.7% respecto al entrenamiento de paralelismo de datos, mientras que DiLoCo mostró una degradación del 16.5%.
  • IsoLoCo incluso superó al baseline de paralelismo de datos con conteos bajos de trabajadores ($R \le 4$).
• Tamaño del Modelo: La ventaja de IsoLoCo sobre DiLoCo persiste a través de los tamaños de modelo (178M a 1B), aunque la mejora relativa es ligeramente menor para modelos más grandes (1.2% en 1B frente a ~2.4% en 178M), probablemente porque los modelos más grandes escalan naturalmente mejor con el número de trabajadores.

Impacto de los Pasos Internos

IsoLoCo fue evaluado con intervalos de sincronización crecientes ($H = 30$ a $240$). Consistientemente superó a DiLoCo en todos los intervalos, demostrando robustez ante periodos más largos de divergencia local.

Comparación con DiLoCo basado en Momento

Notablemente, la combinación de DiLoCo SGD con la agregación Iso-C (sin momento) ya superaba al DiLoCo estándar (que usa momento de Nesterov). Añadir el momento de Nesterov a Iso-C (creando IsoLoCo) amplió aún más la brecha de rendimiento, indicando que la agregación basada en fusión y los mecanismos de momento son complementarios.

Compatibilidad con Muon

El artículo señala que IsoLoCo puede componerse con el optimizador interno Muon (MuLoCo). En esta configuración, IsoLoCo logró la pérdida de validación más baja entre todos los métodos para $R \ge 32$, mientras que IsoLoCo con un optimizador interno AdamW igualó o superó a MuLoCo con conteos altos de trabajadores con costos de ortogonalización significativamente menores.

Significancia y Reivindicaciones

El artículo sostiene que ver el entrenamiento distribuido de baja comunicación a través de la lente de la fusión de modelos proporciona un marco poderoso para mejorar la escalabilidad. Las principales contribuciones son:

1. Unificación Teórica: Establecer una analogía directa entre la agregación de pseudo-gradientes de DiLoCo y la fusión de modelos de aritmética de tareas.
2. Validación Empírica: Demostrar que las técnicas de fusión de modelos, específicamente Iso-C, pueden servir como reglas de agregación superiores para el entrenamiento distribuido, superando el promediado simple e incluso al DiLoCo basado en momento.
3. IsoLoCo: Proponer un algoritmo práctico, potenciado por el momento, que reduce significativamente la brecha de rendimiento entre el entrenamiento de baja comunicación y los estándares de oro de paralelismo de datos, particularmente en regímenes con muchos trabajadores o intervalos de sincronización largos.
4. Eficiencia: Proporcionar una implementación basada en Newton-Schulz que hace que el proceso de fusión, pesado en SVD, sea computacionalmente factible para el entrenamiento distribuido iterativo.

Los autores concluyen que la agregación inspirada en la fusión es una estrategia efectiva para el entrenamiento de baja comunicación, ofreciendo un camino para escalar el entrenamiento a través de clústeres geográficamente distribuidos o mal conectados sin sacrificar la calidad del modelo.