Entrenamiento en multiples GPU

En los últimos años, estamos viendo mejores resultados en muchas tareas comparativas de NLP con modelos de lenguaje preentrenados más grandes . Cómo entrenar redes neuronales grandes y profundas es un desafío, ya que exige una gran cantidad de memoria GPU y un largo horizonte de tiempo de entrenamiento. Sin embargo, un trabajador de GPU individual tiene memoria limitada y los tamaños de muchos modelos grandes han crecido más allá de una sola GPU. Existen varios paradigmas de paralelismo para permitir el entrenamiento de modelos en múltiples GPU, así como una variedad de diseños de arquitectura de modelos y ahorro de memoria para ayudar a hacer posible el entrenamiento de redes neuronales muy grandes .

Paralelismo de entrenamiento

El principal cuello de botella para entrenar modelos de redes neuronales muy grandes es la intensa demanda de una gran cantidad de memoria GPU, muy por encima de lo que se puede alojar en una máquina GPU individual. Además de los pesos del modelo (p. ej., decenas de miles de millones de números de punto flotante), suele ser incluso más caro almacenar resultados de cálculo intermedios, como gradientes y estados del optimizador (p. ej., impulsos y variaciones en Adam). Además, entrenar un modelo grande a menudo se combina con un corpus de entrenamiento grande y, por lo tanto, un solo proceso puede llevar una eternidad. Como resultado, el paralelismo es necesario. El paralelismo puede ocurrir en diferentes dimensiones, incluidos los datos, la arquitectura del modelo y la operación de tensor.

Paralelismo de datos

La forma más ingenua para el paralelismo de datos (DP) es copiar los mismos pesos de modelo en varios trabajadores y asignar una fracción de datos a cada trabajador para que se procese al mismo tiempo. Naive DP no puede funcionar bien si el tamaño del modelo es mayor que la memoria de un solo nodo de GPU. Métodos como GeePS ( Cui et al. 2016 ) descargan parámetros no utilizados temporalmente de vuelta a la CPU para trabajar con memoria de GPU limitada cuando el modelo es demasiado grande para caber en una máquina. La transferencia de intercambio de datos debe ocurrir en el backend y no interferir con el cálculo de entrenamiento.

Al final de cada minilote, los trabajadores deben sincronizar gradientes o pesos para evitar el estancamiento. Hay dos enfoques principales de sincronización y ambos tienen claros pros y contras.

Paralelos síncronos masivos (BSP) : los trabajadores sincronizan datos al final de cada minilote. Evita la obsolescencia de los pesos del modelo y una buena eficiencia de aprendizaje, pero cada máquina tiene que detenerse y esperar a que otras envíen gradientes.
Paralelo asíncrono (ASP) : cada trabajador de la GPU procesa los datos de forma asíncrona, sin esperas ni estancamientos. Sin embargo, puede conducir fácilmente al uso de pesos obsoletos y, por lo tanto, reducir la eficiencia del aprendizaje estadístico. Aunque aumenta el tiempo de cálculo, es posible que no acelere el tiempo de entrenamiento hasta la convergencia.

Modelo de paralelismo

El paralelismo del modelo (MP) tiene como objetivo resolver el caso en el que los pesos del modelo no pueden caber en un solo nodo. Los parámetros de cálculo y modelo se dividen en varias máquinas. A diferencia del paralelismo de datos en el que cada trabajador aloja una copia completa de todo el modelo, MP solo asigna una fracción de los parámetros del modelo a un trabajador y, por lo tanto, se reducen tanto el uso de la memoria como el cálculo.

Dado que las redes neuronales profundas generalmente contienen una pila de capas verticales, parece sencillo dividir un modelo grande por capa, donde un pequeño conjunto consecutivo de capas se agrupa en una partición en un trabajador. Sin embargo, una implementación ingenua para ejecutar cada lote de datos a través de múltiples trabajadores de este tipo con dependencia secuencial genera grandes burbujas de tiempo de espera y una grave infrautilización de los recursos informáticos.

Una configuración de paralelismo de modelo ingenua donde el modelo se divide verticalmente en 4 particiones.

Paralelismo de tuberías

El paralelismo de canalización (PP) combina el paralelismo de modelos con el paralelismo de datos para reducir las "burbujas" de tiempo ineficientes. La idea principal es dividir un minilote en múltiples microlotes y permitir que cada trabajador de etapa procese un microlote simultáneamente. Tenga en cuenta que cada microlote necesita dos pases, uno hacia adelante y otro hacia atrás. La comunicación entre trabajadores solo transfiere activaciones (hacia adelante) y gradientes (hacia atrás). La forma en que se programan estos pases y cómo se agregan los gradientes varía según los diferentes enfoques. El número de particiones (trabajadores) también se conoce como profundidad de tubería .

Paralelismo tensorial

Tanto los paralelismos del modelo como los de la canalización dividen un modelo verticalmente. OTOH, podemos particionar horizontalmente el cálculo para una operación de tensor en múltiples dispositivos, denominado paralelismo de tensor (TP) .

Tomemos el transformador como ejemplo dada su popularidad. El modelo de transformador consiste principalmente en capas de MLP y bloques de autoatención. Megatron-LM ( Shoeybi et al. 2020 ) adopta una forma sencilla de paralelizar el cálculo intracapa para MLP y autoatención.

combinaron canalización, tensor y paralelismo de datos con una nueva estrategia de programación de canalización y llamaron a su enfoque PTD-P . En lugar de colocar solo un conjunto continuo de capas ("fragmento de modelo") en un dispositivo, a cada trabajador se le pueden asignar múltiples fragmentos de subconjuntos de capas continuos más pequeños (por ejemplo, el dispositivo 1 tiene capas 1, 2, 9, 10; el dispositivo 2 tiene capas 3, 4, 11, 12; cada una tiene dos fragmentos modelo).

Mezcla de expertos (MoE)

El enfoque Mixture-of-Experts (MoE) atrae mucha atención recientemente, ya que los investigadores (principalmente de Google) intentan superar el límite del tamaño del modelo. El núcleo de la idea es el aprendizaje por conjuntos : ¡la combinación de varios alumnos débiles le da un alumno fuerte!

Dentro de una red neuronal profunda, el ensamblaje se puede implementar con un mecanismo de activación que conecta a múltiples expertos ( Shazeer et al., 2017 ). El mecanismo de activación controla qué subconjunto de la red (por ejemplo, qué expertos) debe activarse para producir salidas. El documento lo denominó capa de "mezcla de expertos escasamente controlada" (MoE).

Ver más: Mixture-of-Experts (MoE)

Edficiencia memoria

Descarga de CPU

Cuando la memoria de la GPU está llena, una opción es descargar temporalmente los datos no utilizados a la CPU y volver a leerlos cuando sea necesario más tarde ( Rhu et al. 2016 ). La idea de la descarga de la CPU es sencilla, pero es menos popular en los últimos años debido a la ralentización que genera en el tiempo de entrenamiento.

Recálculo de activación

El recálculo de activación (también conocido como "puntos de control de activación" o "puntos de control de gradiente"; Chen et al. 2016 ) es una idea inteligente pero simple para reducir el consumo de memoria a costa del tiempo de cálculo. Reduce el costo de memoria de entrenar un capa de red neuronal profunda

Entrenamiento mixto de precisión

Narang y Micikevicius et al. (2018) introdujeron un método para entrenar modelos utilizando números de punto flotante de precisión media (FP16) sin perder la precisión del modelo.

Tres técnicas para evitar perder información crítica a media precisión:

Copia maestra de pesas de precisión completa . Mantenga una copia de precisión completa (FP32) de los pesos del modelo que acumula gradientes. Los números se redondean a la mitad de precisión para pases hacia adelante y hacia atrás. La motivación es que cada actualización de gradiente (es decir, gradiente multiplicado por la tasa de aprendizaje) puede ser demasiado pequeña para estar completamente contenida dentro del rango FP16 (es decir, se convierte en cero en FP16).
Escalado de pérdidas . Aumente la pérdida para manejar mejor los gradientes con magnitudes pequeñas (consulte la Fig. 16). Escalar los gradientes ayuda a cambiarlos para ocupar una sección más grande hacia la sección derecha (que contiene valores más grandes) del rango representable, conservando los valores que de otro modo se perderían.
Precisión aritmética . Para la aritmética de red común (p. ej., producto punto vectorial, reducción sumando elementos vectoriales), podemos acumular los resultados parciales en FP32 y luego guardar el resultado final como FP16 antes de guardarlo en la memoria. Las operaciones puntuales se pueden ejecutar en FP16 o FP32.

Compresión

Los resultados intermedios a menudo consumen mucha memoria, aunque solo se necesitan en un paso hacia adelante y un paso hacia atrás. Hay una brecha temporal notable entre estos dos usos. Así, Jain et al. (2018) propusieron una estrategia de codificación de datos para comprimir los resultados intermedios después del primer uso en el primer paso y luego decodificarlos para la propagación hacia atrás más tarde.

Su sistema Gist incorpora dos esquemas de codificación: codificación sin pérdidas específica de capa ; concéntrese en los patrones ReLU-Pool ("Binarize") y ReLU-Conv ("Almacenamiento escaso y computación densa"). Codificación con pérdida agresiva ; utilizar la reducción de precisión retardada (DPR). Observaron que el primer uso inmediato de mapas de características debe mantenerse con alta precisión, pero el segundo uso puede tolerar una precisión más baja. Los experimentos demostraron que Gist puede reducir el costo de la memoria 2 veces en 5 DNN de clasificación de imágenes SOTA, con un promedio de 1,8 veces con solo un 4 % de sobrecarga de rendimiento.

Optimizador de memoria eficiente

Los optimizadores están ansiosos por el consumo de memoria. Tome el popular optimizador Adam como ejemplo, internamente necesita mantener impulsos y variaciones, ambos a la misma escala que los gradientes y los parámetros del modelo. De repente, necesitamos ahorrar 4 veces la memoria de los pesos de los modelos.

Se han propuesto varios optimizadores para reducir el consumo de memoria. Por ejemplo, en lugar de almacenar los impulsos completos y las variaciones como en Adam, Adafactor ( Shazeer et al. 2018 ) solo rastrea las sumas por fila y por columna de los promedios móviles y luego estima los segundos momentos en función de estas sumas. SM3 ( Anil et al. 2019 ) describe un método de optimización adaptativo diferente, que también conduce a una memoria muy reducida.

ZeRO ( Zero Redundancy Optimizer ; Rajbhandari et al. 2019 ) optimiza la memoria utilizada para entrenar modelos grandes en función de la observación sobre dos consumos principales de memoria del entrenamiento de modelos grandes:

La mayoría está ocupada por estados del modelo , incluidos los estados del optimizador (p. ej., impulsos y varianzas de Adam), gradientes y parámetros. El entrenamiento de precisión mixta demanda mucha memoria ya que el optimizador necesita mantener una copia de los parámetros del FP32 y otros estados del optimizador, además de la versión del FP16.
El resto es consumido por activaciones, búferes temporales y memoria fragmentada inutilizable (denominados estados residuales en el artículo).

ZeRO combina dos enfoques, ZeRO-DP y ZeRO-R . ZeRO-DP es un paralelismo de datos mejorado para evitar la redundancia simple sobre los estados del modelo. Divide el estado, los gradientes y los parámetros del optimizador en múltiples procesos paralelos de datos a través de un cronograma de comunicación dinámico para minimizar el volumen de comunicación. ZeRO-R optimiza el consumo de memoria de los estados residuales mediante el recálculo de activación dividida, tamaño de búfer constante y desfragmentación de memoria sobre la marcha.

Bibliografía

[1] Gist: Efficient Data Encoding for Deep Neural Network Training

[2] https://lilianweng.github.io/posts/2021-09-25-train-large/.