Entrenar modelos de IA de gran tamaño (como los transformadores y los modelos de lenguaje) se ha convertido en un paso crucial e indispensable en el campo de la IA, pero también enfrenta altos costos computacionales, consumo de memoria y demanda de energía. Por ejemplo, GPT-3 de OpenAI tiene 175 mil millones de parámetros y requiere semanas de entrenamiento en GPU. Esta enorme demanda de recursos limita la aplicación de esta tecnología a organizaciones con acceso a recursos computacionales a gran escala, y también aumenta las preocupaciones sobre la eficiencia energética y el impacto ambiental. Abordar estos desafíos es fundamental para garantizar una accesibilidad y sostenibilidad más amplias del desarrollo de la IA.
Los métodos de entrenamiento tradicionales son ineficientes y requieren soluciones innovadoras.
La principal razón de la baja eficiencia en el entrenamiento de modelos grandes radica en su dependencia de matrices densas, lo que requiere una gran cantidad de memoria y capacidad de cálculo. La compatibilidad limitada de las GPU modernas con operaciones de baja precisión o bajo rango optimizadas exacerba aún más estas necesidades. Si bien se han propuesto algunos métodos, como la descomposición de matrices y la reducción de rango heurística, para mitigar estos problemas, siguen estando limitados en la práctica. Por ejemplo, GaLore puede entrenar en una configuración de lote único, pero tiene una sobrecarga de tiempo de ejecución poco práctica. Del mismo modo, LTE, que utiliza adaptadores de bajo rango, tiene problemas de convergencia en tareas de gran tamaño. Actualmente, falta un método que pueda reducir simultáneamente el uso de memoria, el costo computacional y el tiempo de entrenamiento sin afectar el rendimiento, lo que hace urgente la necesidad de soluciones innovadoras.
Marco CoMERA: entrenamiento eficiente a través de la optimización de tensores adaptativos
Investigadores de la Universidad de Albany (SUNY), la Universidad de California en Santa Bárbara, Amazon Alexa AI y Meta han presentado un nuevo marco llamado CoMERA (método de entrenamiento eficiente en computación y memoria a través de la optimización de tensores de rango adaptativo). Este marco combina eficiencia de memoria y velocidad de cálculo mediante técnicas de compresión de tensores de rango adaptativo. A diferencia de los métodos tradicionales que solo se centran en la compresión, CoMERA utiliza un método de optimización multiobjetivo para equilibrar la relación de compresión y la precisión del modelo. Utiliza incrustaciones tensorizadas y contracciones avanzadas de redes tensoriales para optimizar el uso de la GPU, reduciendo así la sobrecarga de tiempo de ejecución mientras mantiene un rendimiento sólido. El marco también introduce gráficos CUDA para minimizar el retraso en el inicio del núcleo durante las operaciones de la GPU, un cuello de botella importante en los métodos tradicionales de compresión de tensores.
CoMERA se basa en una representación tensorial adaptativa que permite que las capas del modelo ajusten dinámicamente su rango según las restricciones de recursos. Al modificar el rango del tensor, el marco puede lograr la compresión sin comprometer la integridad de las operaciones de la red neuronal. Esta optimización dinámica se logra mediante un proceso de entrenamiento de dos etapas:
Etapa temprana: se centra en la convergencia estable.
Etapa posterior: ajuste fino del rango para cumplir con objetivos de compresión específicos.
En un modelo Transformer de seis codificadores, CoMERA logró una relación de compresión de hasta 43 veces en su etapa temprana, y hasta 361 veces en su etapa de optimización posterior. Además, en comparación con GaLore, redujo el consumo de memoria en 9 veces y aumentó la velocidad de entrenamiento por iteración en 2-3 veces.
Los resultados de múltiples pruebas muestran el excelente rendimiento de CoMERA
Al aplicarse a un modelo Transformer entrenado en el conjunto de datos MNLI, CoMERA redujo el tamaño del modelo de 256 MB a tan solo 3.2 MB, manteniendo la precisión. En sistemas de recomendación a gran escala como DLRM, CoMERA comprimió el modelo en 99 veces y redujo el uso máximo de memoria en 7 veces. El marco también mostró un excelente rendimiento en el preentrenamiento de CodeBERT (un modelo de lenguaje grande de dominio específico), obteniendo una relación de compresión general de 4.23 veces y logrando una aceleración de 2 veces en algunas etapas de entrenamiento. Estos resultados destacan su capacidad para manejar diversas tareas y arquitecturas, ampliando su aplicabilidad en varios campos.
Resumen de las ventajas clave del marco CoMERA
Las principales conclusiones de esta investigación son las siguientes:
CoMERA logró una relación de compresión de hasta 361 veces para capas específicas y de 99 veces para el modelo completo, reduciendo significativamente las necesidades de almacenamiento y memoria.
El marco redujo el tiempo de entrenamiento por iteración de los transformadores y los sistemas de recomendación en 2-3 veces, ahorrando recursos y tiempo de cálculo.
Al utilizar representaciones tensorizadas y gráficos CUDA, CoMERA redujo el consumo máximo de memoria en 7 veces, lo que permite el entrenamiento en GPU más pequeñas.
El método CoMERA admite varias arquitecturas, incluidos los transformadores y los modelos de lenguaje grandes, manteniendo o mejorando la precisión.
Al reducir la energía y los recursos necesarios para el entrenamiento, CoMERA contribuye a prácticas de IA más sostenibles y permite el acceso a modelos de vanguardia a un público más amplio.