Cómo Latent Space utilizó la biblioteca de paralelismo de modelos de Amazon SageMaker para ampliar las fronteras de los transformadores a gran escala

Nodo de origen: 1204406

Este blog es coautor de Sarah Jane Hong CSO, Darryl Barnhart CTO, Ian Thompson, CEO de Latent Space y Prem Ranga de AWS.

El espacio latente es una representación oculta de ideas abstractas que aprenden los modelos de aprendizaje automático (ML). Por ejemplo, "perro", "flor" o "puerta" son conceptos o ubicaciones en el espacio latente. A Espacio latente, estamos trabajando en un motor que le permite manipular y explorar este espacio con indicaciones visuales y de lenguaje. El equipo de Latent Space proviene de dos campos que durante mucho tiempo han tenido poca superposición: gráficos y procesamiento del lenguaje natural (NLP). Tradicionalmente, las modalidades de imágenes y texto se han manejado por separado, cada una con su propia historia de ingeniería de características compleja, costosa y frágil. Las tareas de PNL, como la comprensión de documentos o la respuesta a preguntas, generalmente han tenido poco en común con las tareas de visión, como la comprensión o el renderizado de escenas, y usualmente usamos enfoques y modelos muy diferentes para cada tarea. Pero esto está cambiando rápidamente.

Esta fusión de modalidades en un único espacio latente compartido abre una nueva generación de aplicaciones creativas y comerciales, desde los juegos hasta la comprensión de documentos. Pero desbloquear estas nuevas aplicaciones en un solo modelo abre nuevos desafíos de escala, como se destaca en “La amarga lección” de Richard Sutton, y el emocionante trabajo de los últimos años sobre las leyes de escala. Para que esto sea posible, Latent Space está trabajando en investigación de vanguardia para fusionar estas modalidades en un solo modelo, pero también para escalar y hacerlo de manera eficiente. Aquí es donde entra en juego el paralelismo del modelo.

Amazon SageMakerEl exclusivo enfoque automatizado de partición de modelos y canalización eficiente hizo posible nuestra adopción del paralelismo de modelos con poco esfuerzo de ingeniería, y escalamos nuestro entrenamiento de modelos más allá de mil millones de parámetros (utilizamos el instancias p4d.24xlarge A100), que es un requisito importante para nosotros. Además, observamos que al entrenar con una configuración de entrenamiento de 16 nodos y ocho GPU con la biblioteca de paralelismo del modelo SageMaker, registramos una mejora del 38% en la eficiencia en comparación con nuestras ejecuciones de entrenamiento anteriores.

Desafíos de la formación de transformadores a gran escala

En Latent Space, fusionamos el lenguaje y la visión en modelos de transformadores con miles de millones de parámetros para admitir casos de uso "fuera de distribución" de la imaginación de un usuario o que ocurrirían en el mundo real pero no en nuestros datos de entrenamiento. Estamos manejando los desafíos inherentes al escalado a miles de millones de parámetros y más de dos maneras diferentes:

Las técnicas de recuperación de información han sido durante mucho tiempo un componente clave de los motores de búsqueda y las tareas de control de calidad. Recientemente, se ha logrado un progreso emocionante combinando técnicas IR clásicas con transformadores modernos, específicamente para tareas de respuesta a preguntas donde se entrena un modelo junto con un recuperador neuronal que aprende a recuperar documentos relevantes para ayudar a responder preguntas. Para obtener una descripción general, consulte el trabajo reciente de FAIR en Generación aumentada de recuperación: simplificación de la creación de modelos inteligentes de procesamiento del lenguaje natural y Fusión en decodificador, Google Brain REINOy Nvidia's Perro perdiguero neuronal para responder preguntas.

Si bien las técnicas de recuperación aumentada ayudan con los costos y la eficiencia, todavía no podemos ajustar el modelo en una sola GPU para nuestro modelo más grande. Esto significa que necesitamos usar el paralelismo del modelo para entrenarlo. Sin embargo, debido a la naturaleza de nuestra arquitectura de recuperación, diseñar la división de nuestro modelo fue un desafío debido a las interdependencias entre los contextos recuperados a través de las entradas de entrenamiento. Además, incluso si determinamos cómo dividimos nuestro modelo, la introducción del paralelismo del modelo fue una tarea de ingeniería importante que debíamos realizar manualmente a lo largo de nuestro ciclo de vida de investigación y desarrollo.

La biblioteca de paralelismo de modelos de SageMaker

El paralelismo de modelos es el proceso de dividir un modelo entre varios dispositivos o nodos (como instancias equipadas con GPU) y crear una canalización eficiente para entrenar el modelo en estos dispositivos para maximizar la utilización de la GPU. los modelo de biblioteca de paralelismo en SageMaker hace que el paralelismo de modelos sea más accesible al proporcionar una división de modelos automatizada, también conocida como partición automatizada de modelos y programación sofisticada de ejecución de oleoductos. Los algoritmos de división de modelos pueden optimizar la velocidad o el consumo de memoria. La biblioteca utiliza un algoritmo de partición que equilibra la memoria, minimiza la comunicación entre dispositivos y optimiza el rendimiento.

Partición automatizada de modelos

Para nuestro caso de uso de PyTorch, la biblioteca paralela del modelo ejecuta internamente un paso de seguimiento (en el primer paso de entrenamiento) que construye el gráfico del modelo y determina las formas del tensor y los parámetros. Luego construye un árbol, que consiste en el anidado nn.Module objetos en el modelo, así como datos adicionales recopilados a partir del seguimiento, como la cantidad de nn.Parametersy tiempo de ejecución para cada nn.Module.

Luego, la biblioteca atraviesa este árbol desde la raíz y ejecuta un algoritmo de partición que equilibra la carga computacional y el uso de memoria, y minimiza la comunicación entre instancias. Si varios nn.Modules comparten el mismo nn.Parameter, estos módulos se colocan en el mismo dispositivo para evitar mantener múltiples versiones del mismo parámetro. Después de que se toma la decisión de particionamiento, los módulos y pesos asignados se cargan en sus dispositivos.

Programación de ejecución de canalizaciones

Otra característica principal de la biblioteca paralela de modelos distribuidos de SageMaker es carreras canalizadas, que determinan el orden en el que se realizan los cálculos y se procesan los datos en todos los dispositivos durante el entrenamiento del modelo. La canalización se basa en dividir un mini lote en micro lotes, que se introducen en la canalización de entrenamiento uno por uno y siguen un programa de ejecución definido por el tiempo de ejecución de la biblioteca.

La canalización de microbatch asegura que todas las GPU se utilicen por completo, que es algo que tendríamos que construir nosotros mismos, pero con la biblioteca de paralelismo de modelos, esto sucede perfectamente entre bastidores. Por último, podemos usar Amazonas FSx, lo cual es importante para garantizar que nuestras velocidades de lectura sean rápidas dada la cantidad de archivos que se leen durante el entrenamiento de un modelo multimodal con recuperación.

Arquitectura de formación

El siguiente diagrama representa cómo configuramos nuestra arquitectura de entrenamiento. Nuestros principales objetivos eran mejorar la velocidad del entrenamiento y reducir los costos. Los transformadores de imagen y lenguaje que estamos entrenando son muy complejos, con una cantidad significativamente grande de capas y pesos en su interior, que se ejecutan en miles de millones de parámetros, todo lo cual los hace incapaces de caber en la memoria de un solo nodo. Cada nodo lleva un subconjunto del modelo, a través del cual los datos fluyen y las transformaciones se comparten y compilan. Configuramos 16 p4d.24xgrande instancias cada una con ocho GPU que utilizan la siguiente representación de arquitectura:

A medida que ampliamos nuestros modelos, una tendencia común es tener todo almacenado en los pesos de la red. Sin embargo, por motivos prácticos, queremos aumentar nuestros modelos para aprender a buscar contextos relevantes para ayudar con la tarea de renderizado. Esto nos permite mantener bajos nuestros costos de servicio sin comprometer la calidad de la imagen. Usamos un gran modelo de PNL basado en transformadores y, como se mencionó anteriormente, observamos un aumento del 38% en la eficiencia del entrenamiento con la biblioteca de paralelismo del modelo SageMaker, como se muestra a continuación:

  • Necesitamos una reducción total para cada cálculo en el caso del paralelismo a nivel de tensor. Esto toma O (log_2 n) pasos paralelos. Es decir, n máquinas que toman O (n) pasos, para O (n log_2 n) operaciones totales.
  • Para el paralelismo de la tubería, requerimos O (1) pasos paralelos para pasar datos por la tubería
  • Dadas 16 máquinas con ocho GPU, tenemos O (1) costo para tubería paralela y O (log_2 (8)) = O (3) costo para modelo en profundidad paralelo.
  • En este caso, vemos que el costo de la red se reduce a 1/3 al cambiar a la tubería paralela que usamos con el modelo de paralelismo de SageMaker, y el costo total de capacitación se reduce a 1/2 + 1/2 * 1 / log_2 (16 ) = 0.625 del costo original que conduce a una mejora de la eficiencia correspondiente.

En general, cuando la necesidad justifica una capacitación distribuida (problemas con el tamaño del modelo de escala o los datos de capacitación), podemos seguir un conjunto de mejores prácticas para determinar qué enfoque funciona mejor.

Mejores prácticas para la formación distribuida

Basándonos en nuestra experiencia, sugerimos comenzar con un enfoque paralelo de datos distribuidos. Paralelismo de datos distribuidos como el Biblioteca paralela de datos distribuidos de SageMaker resuelve la mayoría de los problemas de red con réplicas de modelos, por lo que debe ajustar los modelos en la menor cantidad de nodos y luego replicar para escalar el tamaño del lote según sea necesario.

Si se queda sin memoria durante el entrenamiento, como hicimos en este escenario, es posible que desee cambiar a un modelo de enfoque paralelo. Sin embargo, considere estas alternativas antes de probar el entrenamiento en paralelo del modelo:

  • En hardware equipado con NVIDIA Tensor Core, utilice entrenamiento de precisión mixta para aumentar la velocidad y reducir el consumo de memoria.
  • Reduzca el tamaño del lote (o reduzca la resolución de la imagen o la longitud de la secuencia NLP, si es posible).

Además, preferimos diseños de modelos que no tienen normalización por lotes como se describe en Reconocimiento de imágenes a gran escala de alto rendimiento sin normalización. Si no se puede evitar, asegúrese de que la normalización por lotes esté sincronizada en todos los dispositivos. Cuando usa entrenamiento distribuido, su lote se divide en GPU, por lo que las estadísticas de lote precisas requieren sincronización en todos los dispositivos. Sin esto, la normalización aumentará el error y, por lo tanto, afectará la convergencia.

Comience con el entrenamiento en paralelo del modelo cuando tenga las siguientes restricciones:

  • Tu modelo no cabe en un solo dispositivo
  • Debido al tamaño de su modelo, se enfrenta a limitaciones al elegir tamaños de lote más grandes, como si los pesos de su modelo ocupan la mayor parte de la memoria de su GPU y se ve obligado a elegir un tamaño de lote más pequeño y subóptimo

Al optimizar el rendimiento, haga lo siguiente:

  • Utilice la canalización para las comunicaciones entre nodos para minimizar la latencia y aumentar el rendimiento
  • Mantenga las tuberías lo más cortas posible para minimizar las burbujas. El número de microbatches debe ajustarse para equilibrar la eficiencia computacional con el tamaño de la burbuja y ser al menos la longitud de la tubería. Si es necesario, puede formar microbatches a nivel de token como se describe en TeraPipe: paralelismo de canalización a nivel de token para entrenar modelos de lenguaje a gran escala

Cuando optimice el costo, utilice las instancias puntuales administradas por SageMaker para la capacitación. Esto puede optimizar el costo de los modelos de entrenamiento hasta en un 90% en las instancias bajo demanda. SageMaker gestiona las interrupciones de Spot en su nombre.

Otros factores a considerar:

  • Dentro de un nodo, cuando hay una interconexión rápida, hay más matices. Si hay una amplia capacidad de red dentro de los nodos, la reorganización de los datos para una computación más óptima puede mostrar un beneficio.
  • Si las activaciones son mucho más grandes que los tensores de peso, un optimizador fragmentado también puede ayudar. Por favor refiérase a Cero para más información.

La siguiente tabla proporciona algunos escenarios de escalado de entrenamiento comunes y cómo puede configurarlos en AWS.

Guión ¿Cuándo aplica? Solución
Escalado de una sola GPU a muchas GPU Cuando la cantidad de datos de entrenamiento o el tamaño del modelo es demasiado grande Cambie a una instancia de múltiples GPU como p3.16xlarge, que tiene ocho GPU, con los datos y el procesamiento divididos en las ocho GPU, y produce una aceleración casi lineal en el tiempo que lleva entrenar su modelo.
Escalado de una sola instancia a varias instancias Cuando la escala necesita extenderse más allá de cambiar el tamaño de la instancia Escale el número de instancias con la función de estimador de SageMaker Python SDK estableciendo su instance_type en p3.16xlarge y instance_count en 2. En lugar de las ocho GPU en un solo p3.16xlarge, tiene 16 GPU en dos instancias idénticas. Considere usar el Biblioteca paralela de datos distribuidos de SageMaker.
Selección de un modelo de enfoque paralelo para el entrenamiento Al encontrar errores de memoria insuficiente durante el entrenamiento Cambie a un modelo de enfoque paralelo utilizando el Biblioteca paralela de modelos distribuidos de SageMaker.
Rendimiento de la red para comunicaciones entre nodos Para entrenamiento distribuido con múltiples instancias (por ejemplo, comunicación entre los nodos en el clúster cuando se realiza una operación AllReduce) Sus instancias deben estar en la misma región y la misma zona de disponibilidad. Cuando usa el SDK de SageMaker Python, esto se maneja por usted. Tus datos de entrenamiento también deben estar en la misma zona de disponibilidad. Considere usar el Biblioteca paralela de datos distribuidos de SageMaker.
GPU, red y almacenamiento optimizados Para necesidades de formación distribuida a gran escala El tipo de instancia p4d.24xlarge se diseñó para un almacenamiento local rápido y un backplane de red rápido con hasta 400 gigabits, y lo recomendamos encarecidamente como la opción de mayor rendimiento para el entrenamiento distribuido.

Conclusión

Con la biblioteca paralela de modelos en SageMaker, obtenemos muchos de los beneficios listos para usar, como la partición automatizada del modelo y la canalización eficiente. En esta publicación, compartimos nuestros desafíos con nuestro caso de uso de ML, nuestras consideraciones sobre diferentes enfoques de capacitación y cómo usamos la biblioteca de paralelismo de modelos de Amazon SageMaker para acelerar nuestra capacitación. Lo mejor de todo es que ahora puede tomar solo unas pocas horas adoptar las mejores prácticas para el paralelismo del modelo y las mejoras de rendimiento que se describen aquí. Si esta publicación te ayuda o te inspira a resolver un problema, ¡nos encantaría saberlo! Comparta sus comentarios y sugerencias.

Referencias

Para obtener más información, consulte lo siguiente:


Acerca de los autores

Prem Ranga es un arquitecto de soluciones empresariales con sede en Atlanta, GA. Es parte de la comunidad de campo técnico de aprendizaje automático y le encanta trabajar con los clientes en su viaje de aprendizaje automático e inteligencia artificial. Prem es un apasionado de la robótica, es un investigador de vehículos autónomos y también construyó Beer Pours controlado por Alexa en Houston y otros lugares.

sarah jane hong es cofundador y director científico de Latent Space. Su experiencia se encuentra en la intersección de la interacción humano-computadora y el aprendizaje automático. Anteriormente, dirigió la investigación de PNL en Sonar (adquirida por Marchex), que sirve a las empresas en el espacio de la IA conversacional. También es una reconocida desarrolladora de AR / VR, y ha recibido premios y becas de Oculus, Mozilla Mixed Reality y Microsoft Hololens.

Darryl Barnhart es cofundador y director de tecnología de Latent Space. Es un desarrollador experimentado con experiencia en aceleración de GPU, gráficos por computadora, datos a gran escala y aprendizaje automático. Otras pasiones incluyen las matemáticas, el desarrollo de juegos y el estudio de la información.

Ian Thompson es el fundador y CEO de Latent Space. Ian es un ingeniero e investigador inspirado en lo "posible adyacente": tecnologías que están a punto de tener un gran impacto en nuestras vidas. Actualmente se enfoca en simplificar y escalar el aprendizaje de representación multimodal para ayudar a construir una IA segura y creativa. Anteriormente ayudó a construir empresas en gráficos / realidad virtual (AltspaceVR, adquirida por Microsoft) y educación / NLP (HSE).

Fuente: https://aws.amazon.com/blogs/machine-learning/how-latent-space-used-the-amazon-sagemaker-model-parallelism-library-to-push-the-frontiers-of-large- transformadores de escala /

Sello de tiempo:

Mas de Blog de aprendizaje automático de AWS