Ejecute PyTorch Lightning y PyTorch DDP nativo en Amazon SageMaker Training, con Amazon Search

Tantos datos, tan poco tiempo. Los expertos en aprendizaje automático (ML), científicos de datos, ingenieros y entusiastas se han encontrado con este problema en todo el mundo. Desde el procesamiento del lenguaje natural hasta la visión por computadora, desde tablas hasta series temporales y todo lo demás, el antiguo problema de optimizar la velocidad cuando se ejecutan datos en tantas GPU como sea posible ha inspirado innumerables soluciones. Hoy, nos complace anunciar características para los desarrolladores de PyTorch que utilizan marcos nativos de código abierto, como Rayo PyTorch y PyTorchDDP, que agilizará su camino a la nube.

Amazon SageMaker es un servicio completamente administrado para ML, y el entrenamiento de modelos de SageMaker es un entorno informático optimizado para entrenamiento de alto rendimiento a escala. El entrenamiento de modelos de SageMaker ofrece una experiencia de entrenamiento remoto con un plano de control continuo para entrenar y reproducir fácilmente modelos ML con alto rendimiento y bajo costo. Estamos encantados de anunciar nuevas funciones en el portafolio de capacitación de SageMaker que hacen que ejecutar PyTorch a escala sea aún más fácil y accesible:

1. PyTorch Lightning ahora se puede integrar a Biblioteca paralela de datos distribuidos de SageMaker con solo una línea de cambio de código.
2. Entrenamiento de modelos de SageMaker ahora es compatible con PyTorch Distributed Data Parallel nativo con backend NCCL, lo que permite a los desarrolladores migrar a SageMaker más fácil que nunca.

En esta publicación, discutimos estas nuevas características y también aprendemos cómo Amazon Search ha ejecutado PyTorch Lightning con el backend de entrenamiento distribuido optimizado en SageMaker para acelerar el tiempo de entrenamiento del modelo.

Antes de profundizar en el estudio de caso de la Búsqueda de Amazon, para aquellos que no están familiarizados, nos gustaría brindar algunos antecedentes sobre Biblioteca paralela de datos distribuidos de SageMaker. En 2020, desarrollamos y lanzamos una configuración de clúster personalizada para el descenso de gradiente distribuido a escala que aumenta la eficiencia general del clúster, presentada en Amazon Science como Arenque. Con lo mejor de los servidores de parámetros y las topologías basadas en anillos, SageMaker Distributed Data Parallel (SMDDP) está optimizado para la Nube informática elástica de Amazon (Amazon EC2) topología de red, que incluye Educación para Todos. Para tamaños de clúster más grandes, SMDDP puede ofrecer mejoras de rendimiento del 20 al 40 % en relación con Horovod (TensorFlow) y PyTorch Distributed Data Parallel. Para tamaños de clúster más pequeños y modelos compatibles, recomendamos el Compilador de capacitación de SageMaker, que puede reducir el tiempo total de trabajo hasta en un 50 %.

Cliente destacado: PyTorch Lightning en el backend optimizado de SageMaker con Amazon Search

Amazon Search es responsable de la experiencia de búsqueda y descubrimiento en Amazon.com. Potencia la experiencia de búsqueda para los clientes que buscan productos para comprar en Amazon. A un alto nivel, Amazon Search crea un índice para todos los productos vendidos en Amazon.com. Cuando un cliente ingresa una consulta, Amazon Search utiliza una variedad de técnicas de ML, incluidos modelos de aprendizaje profundo, para hacer coincidir productos relevantes e interesantes con la consulta del cliente. Luego clasifica los productos antes de mostrar los resultados al cliente.

Los científicos de Amazon Search han utilizado PyTorch Lightning como uno de los marcos principales para entrenar los modelos de aprendizaje profundo que impulsan la clasificación de búsqueda debido a sus funciones de usabilidad adicionales además de PyTorch. SMDDP no era compatible con los modelos de aprendizaje profundo escritos en PyTorch Lightning antes de este nuevo lanzamiento de SageMaker. Esto evitó que los científicos de Amazon Search que prefieren usar PyTorch Lightning escalaran su entrenamiento de modelos usando técnicas de datos paralelos, lo que ralentizó significativamente su tiempo de entrenamiento y les impidió probar nuevos experimentos que requieren un entrenamiento más escalable.

Los primeros resultados de evaluación comparativa del equipo muestran un tiempo de entrenamiento 7.3 veces más rápido para un modelo de muestra cuando se entrena en ocho nodos en comparación con una línea base de entrenamiento de un solo nodo. El modelo de referencia utilizado en estas evaluaciones comparativas es una red neuronal de perceptrón multicapa con siete capas densas completamente conectadas y más de 200 parámetros. La siguiente tabla resume el resultado de la evaluación comparativa en instancias de capacitación de SageMaker ml.p3.16xlarge.

A continuación, nos sumergimos en los detalles de los nuevos lanzamientos. Si lo desea, puede pasar por nuestro correspondiente cuaderno de ejemplo.

Ejecute PyTorch Lightning con la biblioteca de capacitación distribuida de SageMaker

Nos complace anunciar que SageMaker Data Parallel ahora se integra a la perfección con PyTorch Lightning dentro de la capacitación de SageMaker.

PyTorch Lightning es un marco de código abierto que proporciona una simplificación para escribir modelos personalizados en PyTorch. De alguna manera similar a lo que hizo Keras para TensorFlow, o incluso podría decirse que Hugging Face, PyTorch Lightning proporciona una API de alto nivel con abstracciones para gran parte de la funcionalidad de nivel inferior de PyTorch. Esto incluye la definición del modelo, la creación de perfiles, la evaluación, la poda, el paralelismo del modelo, las configuraciones de hiperparámetros, el aprendizaje de transferencia y más.

Anteriormente, los desarrolladores de PyTorch Lightning no estaban seguros de cómo migrar sin problemas su código de entrenamiento a los clústeres de GPU de SageMaker de alto rendimiento. Además, no tenían forma de aprovechar las ganancias de eficiencia introducidas por SageMaker Data Parallel.

Para PyTorch Lightning, en términos generales, debe haber pocos o ningún cambio en el código para simplemente ejecutar estas API en SageMaker Training. En los cuadernos de ejemplo usamos el DDPEstrategia y Complemento DDPP métodos.

Hay tres pasos para usar PyTorch Lightning con SageMaker Data Parallel como un backend optimizado:

1. Utilice un soporte Contenedor de aprendizaje profundo de AWS (DLC) como su imagen base, u opcionalmente cree su propio contenedor y instale el backend de SageMaker Data Parallel usted mismo. Asegúrese de tener PyTorch Lightning incluido en sus paquetes necesarios, como con un requirements.txt archivo.
2. Realice algunos cambios de código menores en su secuencia de comandos de entrenamiento que habiliten el backend optimizado. Éstos incluyen:
   1. Importe la biblioteca SM DDP:

      import smdistributed.dataparallel.torch.torch_smddp
      

   2. Configure el entorno PyTorch Lightning para SageMaker:

      from pytorch_lightning.plugins.environments.lightning_environment 
        import LightningEnvironment
      
      env = LightningEnvironment()
      env.world_size = lambda: int(os.environ["WORLD_SIZE"])
      env.global_rank = lambda: int(os.environ["RANK"])

   3. Si está utilizando una versión de PyTorch Lightning anterior a la 1.5.10, deberá agregar algunos pasos más.
      1. Primero, agregue la variable de entorno:

         os.environ["PL_TORCH_DISTRIBUTED_BACKEND"] = "smddp"

      2. En segundo lugar, asegúrese de utilizar DDPPlugin, más bien que DDPStrategy. Si está utilizando una versión más reciente, que puede configurar fácilmente colocando el requirements.txt existentes source_dir para su trabajo, entonces esto no es necesario. Ver el siguiente código:

         ddp = DDPPlugin(parallel_devices=[torch.device("cuda", d) for d in range(num_gpus)], cluster_environment=env)

   4. Opcionalmente, defina el backend de su grupo de procesos como "smddp" existentes DDPSTrategy objeto. Sin embargo, si está utilizando PyTorch Lightning con el backend PyTorch DDP, que también es compatible, simplemente elimine este `process_group_backend` parámetro. Ver el siguiente código:

      ddp = DDPStrategy(
        cluster_environment=env, 
        process_group_backend="smddp", 
        accelerator="gpu")

3. Asegúrese de tener un método de distribución anotado en el estimador, como distribution={"smdistributed":{"dataparallel":{"enabled":True} si está utilizando el backend de Herring, o distribution={"pytorchddp":{"enabled":True}.

• Para obtener una lista completa de los parámetros adecuados en el distribution parámetro, consulte nuestra documentación esta página.

¡Ahora puede iniciar su trabajo de capacitación de SageMaker! Puede iniciar su trabajo de capacitación a través de Python SDK, Boto3, la consola de SageMaker, el Interfaz de línea de comandos de AWS (AWS CLI) y muchos otros métodos. Desde la perspectiva de AWS, esto es un solo comando API: create-training-job. Ya sea que ejecute este comando desde su terminal local, un AWS Lambda función, una Estudio Amazon SageMaker notebook, una canalización de KubeFlow o cualquier otro entorno informático depende completamente de usted.

Tenga en cuenta que la integración entre PyTorch Lightning y SageMaker Data Parallel actualmente solo es compatible con las versiones más recientes de PyTorch, a partir de la 1.11. Además, esta versión solo está disponible en los DLC de AWS para SageMaker a partir de PyTorch 1.12. Asegúrate de señalar esta imagen como tu base. En us-east-1, esta dirección es la siguiente:

ecr_image = '763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:1.12.0-gpu-py38-cu113-ubuntu20.04-sagemaker'

Luego puede extender su contenedor Docker usando esto como su imagen base, o puede pasar esto como una variable al image_uri argumento de la Estimador de entrenamiento de SageMaker.

Como resultado, podrá ejecutar su código PyTorch Lightning en las GPU optimizadas de SageMaker Training, con el mejor rendimiento disponible en AWS.

Ejecute datos distribuidos de PyTorch en paralelo en SageMaker

El gran problema Paralelo de datos distribuidos de PyTorch (DDP) resuelve es engañosamente simple: velocidad. Un buen marco de entrenamiento distribuido debe proporcionar estabilidad, confiabilidad y, lo que es más importante, excelente rendimiento a escala. PyTorch DDP cumple con esto al proporcionar a los desarrolladores de antorcha API para replicar sus modelos en múltiples dispositivos de GPU, tanto en configuraciones de un solo nodo como de múltiples nodos. Luego, el marco administra la fragmentación de diferentes objetos del conjunto de datos de entrenamiento en cada copia del modelo, promediando los gradientes para cada una de las copias del modelo para sincronizarlos en cada paso. Esto produce un modelo en la finalización total de la ejecución de entrenamiento completa. El siguiente diagrama ilustra este proceso.

PyTorch DDP es común en proyectos que utilizan grandes conjuntos de datos. El tamaño preciso de cada conjunto de datos variará ampliamente, pero una pauta general es escalar conjuntos de datos, calcular tamaños y modelar tamaños en proporciones similares. También llamado leyes de escala, la combinación óptima de estos tres es objeto de debate y variará según las aplicaciones. En AWS, basándonos en el trabajo con varios clientes, podemos ver claramente los beneficios de las estrategias de datos paralelos cuando el tamaño general de un conjunto de datos es de al menos algunas decenas de GB. Cuando los conjuntos de datos aumentan aún más, la implementación de algún tipo de estrategia de datos paralelos es una técnica fundamental para acelerar el experimento general y mejorar el tiempo de generación de valor.

Anteriormente, los clientes que usaban PyTorch DDP para capacitación distribuida en las instalaciones o en otros entornos informáticos carecían de un marco para migrar fácilmente sus proyectos a SageMaker Training para aprovechar las GPU de alto rendimiento con un plano de control continuo. Específicamente, necesitaban migrar su marco paralelo de datos a SMDDP o desarrollar y probar las capacidades de PyTorch DDP en SageMaker Training manualmente. Hoy, SageMaker Training se complace en brindar una experiencia perfecta para los clientes que incorporan su código PyTorch DDP.

Para usar esto de manera efectiva, no necesita hacer ningún cambio en sus scripts de entrenamiento.

Puede ver este nuevo parámetro en el siguiente código. En el distribution parámetro, simplemente agregue pytorchddp y establecer habilitado como true.

estimator = PyTorch(
    base_job_name="pytorch-dataparallel-mnist",
    source_dir="code",
    entry_point = "my_model.py",
    ... 
    # Training using SMDataParallel Distributed Training Framework
    distribution = {"pytorchddp": {"enabled": "true"}}
)

Esta nueva configuración comienza en SageMaker Python SDK versiones 2.102.0 y PyTorch DLC 1.11.

Para los desarrolladores de PyTorch DDP que están familiarizados con el popular carrera de antorchas Framework, es útil saber que esto no es necesario en el entorno de capacitación de SageMaker, que ya proporciona una sólida tolerancia a fallas. Sin embargo, para minimizar las reescrituras de código, puede traer otro script de inicio que ejecute este comando como su punto de entrada.

Ahora los desarrolladores de PyTorch pueden mover fácilmente sus scripts a SageMaker, lo que garantiza que sus scripts y contenedores puedan ejecutarse sin problemas en múltiples entornos informáticos.

Esto los prepara para, en el futuro, aprovechar Bibliotecas de formación distribuidas de SageMaker que proporcionan topologías de entrenamiento optimizadas para AWS para ofrecer hasta un 40 % de mejoras en la aceleración. Para los desarrolladores de PyTorch, ¡esto es una sola línea de código! Para el código PyTorch DDP, simplemente puede configurar el backend para smddp en la inicialización (ver Modificar un script de entrenamiento de PyTorch), como se muestra en el siguiente código:

import smdistributed.dataparallel.torch.torch_smddp
import torch.distributed as dist
dist.init_process_group(backend='smddp')

Como vimos anteriormente, también puede configurar el backend de DDPStrategy a smddp al usar Rayo. Esto puede conducir a hasta 40% de aceleraciones generales para grandes racimos! Para obtener más información sobre la capacitación distribuida en SageMaker, consulte nuestro webinar a pedidoapoyando ordenadores portátiles, relevante documentacióny papeles.

Conclusión

En esta publicación, presentamos dos nuevas funciones dentro de la familia de SageMaker Training. Esto hace que sea mucho más fácil para los desarrolladores de PyTorch usar su código existente en SageMaker, tanto PyTorch DDP como PyTorch Lightning.

También mostramos cómo Amazon Search usa SageMaker Training para entrenar sus modelos de aprendizaje profundo y, en particular, PyTorch Lightning con la biblioteca colectiva optimizada de SageMaker Data Parallel como backend. Pasar a la capacitación distribuida en general ayudó a Amazon Search a lograr tiempos de tren 7.3 veces más rápidos.

Sobre los autores

emilio webber se unió a AWS justo después del lanzamiento de SageMaker y ha estado tratando de contárselo al mundo desde entonces. Además de crear nuevas experiencias de aprendizaje automático para los clientes, a Emily le gusta meditar y estudiar el budismo tibetano.

Karan Dhiman es ingeniero de desarrollo de software en AWS, con sede en Toronto, Canadá. Le apasiona el espacio de aprendizaje automático y la creación de soluciones para acelerar las cargas de trabajo informáticas distribuidas.

Vishwa Karia es ingeniero de desarrollo de software en AWS Deep Engine. Sus intereses se encuentran en la intersección de Machine Learning y Distributed Systems y también le apasiona empoderar a las mujeres en tecnología e inteligencia artificial.

Eiman Elnahrawy es un ingeniero de software principal en Amazon Search que lidera los esfuerzos de aceleración, escalado y automatización del aprendizaje automático. Su experiencia abarca múltiples áreas, incluido el aprendizaje automático, los sistemas distribuidos y la personalización.