Se presenta PyTorch 2.0

22 de marzo de 2023

PyTorch es una biblioteca de aprendizaje automático de código abierto basada en Python que permite a los desarrolladores crear y entrenar modelos de aprendizaje profundo de manera eficiente.

La nueva generación de la popular biblioteca de aprendizaje profundo

PyTorch 2.0 marca el inicio de la próxima generación de la serie 2 de PyTorch. A lo largo de los últimos años, se ha innovado e iterado desde PyTorch 1.0 hasta la versión más reciente, 1.13, y se ha trasladado a la recién formada PyTorch Foundation, parte de la Linux Foundation. La mayor fortaleza de PyTorch, además de su increíble comunidad, es su integración de primera clase con Python, el estilo imperativo, la simplicidad de la API y las opciones que ofrece. PyTorch 2.0 mantiene la misma experiencia de desarrollo en modo ávido, pero cambia y potencia fundamentalmente cómo PyTorch opera a nivel de compilador. Esto permite un rendimiento más rápido y soporte para formas dinámicas y distribuidas.

PyTorch 2.x: más rápido, más «pythonic» y tan dinámico como siempre

La nueva característica torch.compile lleva el rendimiento de PyTorch a nuevas alturas y comienza el proceso de trasladar partes de PyTorch de C++ a Python. Esta es una dirección sustancialmente nueva para PyTorch, por lo que se le denomina 2.0. torch.compile es una característica completamente aditiva (y opcional), lo que significa que la versión 2.0 es 100% compatible con versiones anteriores.

Nuevas tecnologías en PyTorch 2.0

Las tecnologías subyacentes en torch.compile incluyen TorchDynamo, AOTAutograd, PrimTorch y TorchInductor:

TorchDynamo captura programas de PyTorch de manera segura utilizando Python Frame Evaluation Hooks y es una innovación significativa que resultó de 5 años de investigación y desarrollo en la captura segura de gráficos.
AOTAutograd sobrecarga el motor de autograd de PyTorch como un autodiff de trazado para generar trazas hacia atrás antes de tiempo.
PrimTorch canonicaliza más de 2000 operadores de PyTorch en un conjunto cerrado de aproximadamente 250 operadores primitivos que los desarrolladores pueden utilizar para construir un backend completo de PyTorch. Esto reduce sustancialmente la barrera para escribir una característica o backend de PyTorch.
TorchInductor es un compilador de aprendizaje profundo que genera código rápido para múltiples aceleradores y backends. Para GPUs NVIDIA, utiliza OpenAI Triton como un bloque de construcción clave.

Estas tecnologías están escritas en Python y admiten formas dinámicas (es decir, la capacidad de enviar tensores de diferentes tamaños sin inducir una recompilación), lo que las hace flexibles, fácilmente modificables y reduce la barrera de entrada para desarrolladores y proveedores.

Validación y resultados

Para validar estas tecnologías, se utilizó un conjunto diverso de 163 modelos de código abierto en varios dominios de aprendizaje automático. Se incluyeron tareas como clasificación de imágenes, detección de objetos, generación de imágenes, varias tareas de PNL como modelado de lenguaje, preguntas y respuestas, clasificación de secuencias, sistemas de recomendación y aprendizaje por refuerzo. Los modelos se dividieron en tres categorías:

46 modelos de HuggingFace Transformers
61 modelos de TIMM: una colección de modelos de imágenes de PyTorch de última generación de Ross Wightman
56 modelos de TorchBench: un conjunto seleccionado de bases de código populares de GitHub

Los modelos de código abierto no se modificaron, excepto para agregar una llamada torch.compile que los envuelve. Luego, se midieron las mejoras en velocidad y se validó la precisión en estos modelos. Se informó una velocidad promedio ponderada desigual de 0.75 * AMP + 0.25 * float32, ya que se encontró que AMP es más común en la práctica.

En estos 163 modelos de código abierto, torch.compile funciona el 93% del tiempo y el modelo se ejecuta un 43% más rápido en el entrenamiento en una GPU NVIDIA A100. Con precisión Float32, funciona un 21% más rápido en promedio y con precisión AMP, funciona un 51% más rápido en promedio.

Advertencias

En una GPU de clase de escritorio como NVIDIA 3090, se ha medido que las mejoras en velocidad son menores que en las GPU de nivel empresarial como A100.

Cómo empezar

Para obtener más información sobre PyTorch 2.0 y cómo comenzar a usarlo, consulte los tutoriales, requisitos, modelos y preguntas frecuentes disponibles en el sitio web oficial de PyTorch.