YaFSDP, una herramienta de formación LLM que reduce el uso de GPU en un 20 %, ya está disponible

El desarrollo de grandes modelos de lenguaje requiere inversiones sustanciales en tiempo y recursos de GPU, lo que se traduce directamente en altos costos. Cuanto más grande es el modelo, más pronunciados se vuelven estos desafíos.

Recientemente, Yandex ha introducido una nueva solución: YaFSDP , una herramienta de código abierto que promete revolucionar la formación LLM al reducir significativamente el consumo de recursos de GPU y el tiempo de formación. En un escenario de preentrenamiento que involucra un modelo con 70 mil millones de parámetros, el uso de YaFSDP puede ahorrar los recursos de aproximadamente 150 GPU. Esto se traduce en ahorros mensuales potenciales de aproximadamente entre 0,5 y 1,5 millones de dólares, según la plataforma o el proveedor de GPU virtual.

Yandex ha puesto YaFSDP a disposición del público en GitHub .

El desafío de la formación LLM distribuida

La capacitación de LLM en múltiples GPU implica operaciones complejas que generan ineficiencias y un alto consumo de memoria. Uno de los principales problemas es la necesidad de enviar y recibir cantidades masivas de datos entre GPU. Por ejemplo, en una operación típica all_reduce, se debe comunicar el doble de datos de gradiente que parámetros de red. En el caso del modelo Llama 70B, esto supone transferir 280 GB de datos por iteración.

Además, los pesos, gradientes y estados del optimizador se duplican en las GPU, lo que genera una enorme carga de memoria. El modelo Llama 70B y el optimizador Adam requieren más de 1 TB de memoria, lo que supera con creces la capacidad de memoria típica de 80 GB de la mayoría de las GPU. Esta redundancia ralentiza gravemente el proceso de entrenamiento y, a menudo, hace que sea poco práctico colocar incluso modelos de tamaño moderado en la memoria de la GPU.

Presentamos YaFSDP

YaFSDP de Yandex ofrece una solución muy eficaz a estos desafíos. Al centrarse en optimizar el consumo de memoria y eliminar los cuellos de botella en la comunicación, YaFSDP mejora la eficiencia de la formación LLM. Funciona fragmentando capas en lugar de parámetros individuales, manteniendo comunicaciones eficientes y evitando operaciones redundantes. Además, YaFSDP preasigna buffers para todos los datos requeridos, asegurando que el asignador de Torch no introduzca ineficiencias.

YaFSDP opera utilizando dos buffers para pesos y gradientes intermedios, con capas impares usando un buffer y capas pares usando el otro.

Los pesos de diferentes capas se almacenan en la misma memoria. Si las capas tienen la misma estructura, siempre serán idénticas. Es crucial asegurarse de que cuando necesite la capa X, el búfer contenga los pesos para la capa X. Todos los parámetros se almacenarán en el fragmento de memoria correspondiente dentro del búfer.

Consumo de memoria

Durante el entrenamiento, los principales consumidores de memoria son los pesos, los gradientes, los estados del optimizador, los buffers y las activaciones. YaFSDP reduce significativamente el consumo de memoria al optimizar cómo se almacenan y acceden a estos elementos.

• Pesos, gradientes y estados del optimizador : dependen de la cantidad de procesos y su consumo de memoria tiende a acercarse a cero a medida que aumenta la cantidad de procesos. Al fragmentar estos componentes entre GPU, YaFSDP minimiza la duplicación y, por lo tanto, reduce el uso de memoria.

• Los buffers consumen una cantidad constante de memoria y almacenan valores intermedios durante los cálculos.

• Las activaciones dependen del tamaño del modelo y de la cantidad de tokens procesados por GPU.

Puntos de control de activación

El control de activación es una técnica que almacena solo las activaciones necesarias durante el pase hacia adelante y las recalcula durante el pase hacia atrás. Esto reduce significativamente el uso de memoria, ya que sólo se almacenan los datos esenciales. Por ejemplo, al entrenar un modelo Llama 2 70B con un tamaño de lote de 8192 tokens, el almacenamiento de activación se puede reducir de más de 110 GB a solo 5 GB.

Sin embargo, este enfoque introduce una sobrecarga computacional adicional, que YaFSDP permite evitar al no utilizar los puntos de control de activación para algunas capas, lo cual es posible debido a la optimización de la memoria.

Optimización de la comunicación

YaFSDP mejora la eficiencia de la comunicación de la GPU al garantizar que los datos se transfieran solo cuando sea necesario y al superponer la comunicación con la computación. Utiliza flujos CUDA para gestionar cálculos y comunicaciones simultáneos de forma eficaz.

La herramienta utiliza dos flujos: un flujo de cálculo y un flujo de comunicación. Los eventos sincronizan estos flujos, asegurando que las operaciones se ejecuten en el orden correcto sin introducir puntos muertos.

El paso hacia adelante en la tercera capa no comienza hasta que se completa la operación all_gather (condición 1). Del mismo modo, la operación all_gather en la tercera capa no comenzará hasta que se complete el paso hacia adelante en la primera capa que usa el mismo búfer (condición 2). Como no hay ciclos en este esquema, el punto muerto es imposible.

Resultados experimentales y mejoras de rendimiento

La implementación de YaFSDP ha mostrado mejoras notables en la eficiencia de la capacitación. En un escenario de preentrenamiento con un modelo que tenía 70 mil millones de parámetros, YaFSDP pudo ahorrar los recursos de aproximadamente 150 GPU. Esto se traduce en importantes ahorros de costos mensuales, que oscilan entre $ 0,5 y $ 1,5 millones, según el proveedor o la plataforma de GPU virtual.

YaFSDP reduce el tiempo de entrenamiento hasta en un 26% en comparación con métodos existentes como FSDP y optimiza el uso de la memoria, lo que permite entrenar modelos más grandes de manera más eficiente.

Yandex ha puesto YaFSDP a disposición del público en GitHub . Los ingenieros de ML pueden aprovechar esta herramienta para mejorar la eficiencia de sus procesos de formación de LLM. Al abrir YaFSDP, Yandex pretende fomentar la innovación y la colaboración en la comunidad de IA, permitiendo a los desarrolladores entrenar modelos de forma más rápida y rentable.

YaFSDP representa un avance significativo en la formación LLM. Abordar los desafíos críticos del consumo de memoria y las ineficiencias de la comunicación permite un entrenamiento más rápido y eficiente de grandes modelos de lenguaje.