En la última década, el aprendizaje por refuerzo (RL) se ha convertido en una de las áreas de investigación más prometedoras en el aprendizaje automático y ha demostrado un gran potencial para resolver problemas sofisticados del mundo real, como la ubicación de chips y la gestión de recursos , y la resolución de juegos desafiantes (p. Ej., Ve , Dota 2 y al escondite ). En términos más simples, una infraestructura de RL es un ciclo de recolección de datos y capacitación, donde los actores exploran el entorno y recolectan muestras, que luego se envían a los alumnos para capacitar y actualizar el modelo. La mayoría de las técnicas actuales de RL requieren muchas iteraciones sobre lotes de millones de muestras del entorno para aprender una tarea objetivo (por ejemplo, Dota 2 aprende de lotes de 2 millones de fotogramas cada 2 segundos). Como tal, una infraestructura de RL no solo debe escalar de manera eficiente (por ejemplo, aumentar el número de actores) y recolectar una inmensa cantidad de muestras, sino también ser capaz de iterar rápidamente sobre estas grandes cantidades de muestras durante el entrenamiento.
 |
| Descripción general de un sistema RL en el que un actor envía trayectorias (por ejemplo, múltiples muestras) a un alumno. El alumno entrena un modelo utilizando los datos muestreados y envía el modelo actualizado al actor (por ejemplo, TF-Agents , IMPALA ). |
Hoy presentamos Menger 1 , una infraestructura RL distribuida masiva a gran escala con inferencia localizada que escala hasta varios miles de actores en múltiples clústeres de procesamiento (por ejemplo, células Borg ), lo que reduce el tiempo de capacitación general en la tarea de colocación de chips. En esta publicación, describimos cómo implementamos Menger utilizando aceleradores de TPU de Google para iteraciones de entrenamiento rápidas y presentamos su rendimiento y escalabilidad en la desafiante tarea de la colocación de chips . Menger reduce el tiempo de formación hasta 8,6 x en comparación con una implementación de referencia.
Diseño del sistema Menger
Hay varios sistemas RL distribuidos, como Acme y SEED RL , cada uno de los cuales se enfoca en optimizar un único punto de diseño particular en el espacio de los sistemas de aprendizaje por refuerzo distribuido. Por ejemplo, mientras que Acme utiliza la inferencia local en cada actor con la recuperación frecuente del modelo del alumno, SEED RL se beneficia de un diseño de inferencia centralizado al asignar una parte de los núcleos de TPU para realizar llamadas por lotes. Las compensaciones entre estos puntos de diseño son (1) pagar el costo de comunicación de enviar / recibir observaciones y acciones a / desde un servidor de inferencia centralizado o pagar el costo de comunicación de la recuperación del modelo de un alumno y (2) el costo de la inferencia en los actores ( por ejemplo, CPU) en comparación con aceleradores (por ejemplo, TPU / GPU). Debido a los requisitos de nuestra aplicación objetivo (por ejemplo, tamaño de las observaciones, acciones y tamaño del modelo), Menger usa la inferencia local de una manera similar a Acme, pero lleva la escalabilidad de los actores a un límite virtualmente ilimitado. Los principales desafíos para lograr una escalabilidad masiva y un entrenamiento rápido en aceleradores incluyen:
- Dar servicio a un gran número de solicitudes de lectura de los actores a un alumno para la recuperación del modelo puede estrangular fácilmente al alumno y convertirse rápidamente en un cuello de botella importante (por ejemplo, aumentar significativamente el tiempo de convergencia) a medida que aumenta el número de actores.
- El rendimiento de la TPU a menudo está limitado por la eficiencia de la canalización de entrada en la alimentación de los datos de entrenamiento a los núcleos informáticos de la TPU. A medida que aumenta la cantidad de núcleos de procesamiento de TPU (p. Ej., Pod de TPU ), el rendimiento de la canalización de entrada se vuelve aún más crítico para el tiempo de ejecución de entrenamiento general.
Recuperación eficiente de modelos
Para abordar el primer desafío, presentamos componentes de almacenamiento en caché transparentes y distribuidos entre el alumno y los actores optimizados en TensorFlow y respaldados por Reverb (enfoque similar al utilizado en Dota ). La principal responsabilidad de los componentes de almacenamiento en caché es lograr un equilibrio entre la gran cantidad de solicitudes de los actores y el trabajo del alumno. Agregar estos componentes de almacenamiento en caché no solo reduce significativamente la presión sobre el alumno para atender las solicitudes de lectura, sino que también distribuye aún más a los actores en múltiplesceldas Borg con una sobrecarga de comunicación marginal. En nuestro estudio, mostramos que para un modelo de 16 MB con 512 actores, los componentes de almacenamiento en caché introducidos reducen la latencia de lectura promedio en un factor de ~ 4.0 x, lo que conduce a iteraciones de entrenamiento más rápidas, especialmente para algoritmos de políticas como PPO .
 |
| Descripción general de un sistema RL distribuido con múltiples actores colocados en diferentes células Borg. Dar servicio a las frecuentes solicitudes de actualización del modelo de una gran cantidad de actores en diferentes células Borg acelera al alumno y la red de comunicación entre el alumno y los actores, lo que conduce a un aumento significativo en el tiempo de convergencia general. Las líneas punteadas representan la comunicación gRPC entre diferentes máquinas. |
 |
| Descripción general de un sistema RL distribuido con múltiples actores colocados en diferentes celdas Borg con el servicio de almacenamiento en caché transparente y distribuido introducido. El alumno solo envía el modelo actualizado a los servicios de almacenamiento en caché distribuidos. Cada servicio de almacenamiento en caché maneja las actualizaciones de solicitud de modelo de los actores cercanos (es decir, actores ubicados en las mismas celdas Borg) y el servicio de almacenamiento en caché. El servicio de almacenamiento en caché no solo reduce la carga del alumno para atender las solicitudes de actualización del modelo, sino que también reduce la latencia de lectura promedio por parte de los actores. |
Canal de entrada de alto rendimiento
Para ofrecer una canalización de datos de entrada de alto rendimiento, Menger utiliza Reverb , un sistema de almacenamiento de datos de código abierto recientemente diseñado para aplicaciones de aprendizaje automático que proporciona una plataforma eficiente y flexible para implementar la reproducción de experiencias en una variedad de algoritmos dentro y fuera de las políticas. Sin embargo, el uso de un único servicio de búfer de reproducción de Reverb actualmente no se escala bien en un entorno RL distribuido con miles de actores y simplemente se vuelve ineficiente en términos de rendimiento de escritura de los actores.
 |
| Un sistema RL distribuido con un solo búfer de reproducción. Dar servicio a una gran cantidad de solicitudes de escritura de los actores acelera el búfer de reproducción y reduce su rendimiento general. Además, a medida que escalamos al alumno a un entorno con varios motores de procesamiento (p. Ej., TPU Pod), la alimentación de datos a estos motores desde un único servicio de búfer de reproducción se vuelve ineficiente, lo que afecta negativamente el tiempo de convergencia general. |
Para comprender mejor la eficiencia del búfer de reproducción en una configuración distribuida, evaluamos la latencia de escritura promedio para varios tamaños de carga útil de 16 MB a 512 MB y una cantidad de actores que van desde 16 a 2048. Repetimos el experimento cuando el búfer de reproducción y los actores se colocan en la misma celda Borg . A medida que aumenta el número de actores, la latencia de escritura promedio también aumenta significativamente. Al expandir el número de actores de 16 a 2048, la latencia de escritura promedio aumenta en un factor de ~ 6.2 xy ~ 18.9 x para un tamaño de carga útil de 16 MB y 512 MB, respectivamente. Este aumento en la latencia de escritura afecta negativamente el tiempo de recopilación de datos y conduce a la ineficiencia en el tiempo de formación general.
 |
| La latencia de escritura promedio en un solo búfer de reproducción de Reverb para varios tamaños de carga útil (16 MB - 512 MB) y varios números de actores (16 a 2048) cuando los actores y el búfer de reproducción se colocan en las mismas celdas Borg. |
Para mitigar esto, utilizamos la capacidad de fragmentación proporcionada por Reverb para aumentar el rendimiento entre los actores, el alumno y los servicios de búfer de reproducción. La fragmentación equilibra la carga de escritura de la gran cantidad de actores en varios servidores de búfer de reproducción, en lugar de limitar un solo servidor de búfer de reproducción, y también minimiza la latencia de escritura promedio para cada servidor de búfer de reproducción (ya que menos actores comparten el mismo servidor). Esto permite a Menger escalar de manera eficiente a miles de actores en múltiples células Borg.
 |
| Un sistema RL distribuido con búferes de reproducción fragmentados . Cada servicio de búfer de reproducción es un almacenamiento de datos dedicado para una colección de actores, generalmente ubicados en las mismas celdas Borg. Además, la configuración de búfer de reproducción fragmentada proporciona una canalización de entrada de mayor rendimiento a los núcleos del acelerador. |
Estudio de caso: colocación de chips
Estudiamos los beneficios de Menger en la compleja tarea de colocar chips para una gran lista de redes . Usando 512 núcleos de TPU, Menger logra mejoras significativas en el tiempo de entrenamiento (hasta ~ 8.6 x , reduciendo el tiempo de entrenamiento de ~ 8.6 horas a solo una hora en la configuración más rápida) en comparación con una base sólida. Si bien Menger se optimizó para TPU, el factor clave para esta ganancia de rendimiento es la arquitectura, y esperaríamos ver ganancias similares cuando se adapte para su uso en GPU.
 |
| La mejora en el tiempo de entrenamiento utilizando Menger con número variable de núcleos de TPU en comparación con una línea de base en la tarea de colocación de chips . |
Creemos que la infraestructura de Menger y sus resultados prometedores en la intrincada tarea de la colocación de chips demuestran un camino innovador para acortar aún más el ciclo de diseño de chips y tiene el potencial de no solo permitir más innovaciones en el proceso de diseño de chips, sino otros desafíos del mundo real. tareas también.
Expresiones de gratitud
La mayor parte del trabajo fue realizado por Amir Yazdanbakhsh, Junchao Chen y Yu Zheng. También nos gustaría agradecer a Robert Ormandi, Ebrahim Songhori, Shen Wang, TF-Agents team, Albin Cassirer, Aviral Kumar, James Laudon, John Wilkes, Joe Jiang, Milad Hashemi, Sat Chatterjee, Piotr Stanczyk, Sabela Ramos, Lasse Espeholt, Marcin Michalski, Sam Fishman, Ruoxin Sang, Azalia Mirhosseini, Anna Goldie y Eric Johnson por su ayuda y apoyo.
1 Un cubo de Menger es una curva fractal tridimensional, y la inspiración para el nombre de este sistema, dado que la infraestructura propuesta puede escalar virtualmente ad infinitum. ↩
