PyTorch 2.0 ya está aquí: todo lo que sabemos

PyTorch es un framework de deep learning de código abierto, impulsado por la comunidad, que ofrece una forma flexible y eficiente de crear modelos de machine learning. Cuenta con una interfaz sencilla, un gran apoyo de la comunidad y una integración perfecta con el ecosistema de Python.

PyTorch 2.0 introduce cambios fundamentales en las operaciones del compilador central, manteniendo la misma familiaridad y facilidad de uso para desarrolladores. Esta mejora promete un rendimiento más rápido y mayor compatibilidad con Dynamic Shapes y Distributed.

¿Qué hay de nuevo en PyTorch 2.0?

PyTorch está trasladando componentes de C++ de vuelta a Python, haciéndolo más rápido y fácil de modificar. Con la versión 2.0, han introducido torch.compile, que cambia cómo funciona PyTorch a nivel de compilador. Esta función es opcional y no afecta a tu código anterior.

Compilación en PyTorch 2.0

Para dar una base sólida a la función torch.compile, se han incorporado nuevas tecnologías:

• TorchDynamo. Un compilador JIT (Just-in-Time) a nivel de Python diseñado específicamente para acelerar PyTorch. Al integrarse con la API de evaluación de frames de CPython, modifica dinámicamente el bytecode de Python en tiempo de ejecución, lo que permite ejecutar el código más rápido.
• AOTAutograd. Un conjunto de herramientas que ayuda a acelerar el entrenamiento de modelos en PyTorch. Traza por adelantado los grafos forward y backward. Además, AOTAutograd ofrece mecanismos sencillos para compilar los grafos extraídos con compiladores de deep learning de última generación.
• PrimTorch. Al reducir drásticamente el número de operadores de PyTorch de más de 2000 a un conjunto conciso de aproximadamente 250 operadores primitivos, PrimTorch ha simplificado notablemente el desarrollo de nuevas funciones o backends para PyTorch.
• TorchInductor. Un compilador de deep learning nativo de PyTorch que convierte automáticamente modelos de PyTorch en código generado para múltiples aceleradores y backends. TorchInductor utiliza OpenAI Triton como bloque de construcción para la aceleración en GPU.

Todas estas tecnologías nuevas están escritas en Python y admiten Dynamic shapes. Hacen que el nuevo PyTorch ejecute el código más rápido, sea más flexible y fácil de personalizar, reduciendo la barrera de entrada.

Ejemplos de código

Veamos una implementación rápida y sencilla del compilador de PyTorch.

Sin torch.compile

import torch
model = torch.hub.load("pytorch/vision", "resnet50", weights="IMAGENET1K_V2")

Con torch.compile

Para exprimir el rendimiento del modelo, añade el wrapper torch.compile alrededor del modelo y obtén un modelo compilado. Es plug and play.

import torch
model = torch.hub.load("pytorch/vision", "resnet50", weights="IMAGENET1K_V2")
compiled_model = torch.compile(model)

Aprende a crear modelos de deep learning con la librería PyTorch en nuestro curso Deep Learning with PyTorch.

Puedes entrenar tu modelo sin cambiar nada.

import torch
model = torch.compile(model)
for batch in dataloader:
    run_epoch(model, batch)

O ejecutar inferencia.

model = torch.compile(model)
model(**input)

torch.compile() incluye parámetros adicionales:

1. mode: especifica en qué debe centrarse el compilador al optimizar.
2. dynamic: habilita la ruta de código para Dynamic Shapes.
3. fullgraph: compila el programa en un único grafo.
4. backend: por defecto usa TorchInductor, pero puedes indicar otros backends de compilación disponibles.

def torch.compile(model: Callable,
  *,
  mode: Optional[str] = "default",
  dynamic: bool = False,
  fullgraph:bool = False,
  backend: Union[str, Callable] = "inductor",
  **kwargs
) -> torch._dynamo.NNOptimizedModule

Aprende los fundamentos de la API de PyTorch y crea desde cero una red neuronal sencilla con nuestro tutorial de PyTorch.

Benchmark

Para empezar con PyTorch 2.0, se han utilizado 163 modelos de código abierto (46 de HuggingFace Transformers, 61 de TIMM y 56 de TorchBench) para crear benchmarks de rendimiento de la nueva función de compilación. El benchmark incluye tareas como clasificación de imágenes, generación de imágenes, modelado de lenguaje, sistemas de recomendación y aprendizaje por refuerzo.

Los resultados muestran una mejora de rendimiento significativa durante el entrenamiento en GPUs NVIDIA A100.

Nota: actualmente, el backend por defecto solo es compatible con CPUs y GPUs Nvidia de las series Volta y Ampere.

Benchmark del compilador de PyTorch en GPU NVIDIA A100

Es solo el principio: en próximas actualizaciones verás más mejoras de rendimiento y escalabilidad.

Cómo instalar PyTorch 2.0

Puedes instalar la versión más reciente de PyTorch con pip.

Copia y pega el siguiente comando en tu terminal.

Para GPUs: CUDA 11.8

Las GPUs más recientes han mostrado un rendimiento muy superior.

pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cu118

Para GPUs: CUDA 11.7

pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cu117

Para CPUs:

pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cpu

Para verificación:

git clone https://github.com/pytorch/pytorch
cd tools/dynamo
python verify_dynamo.py

Acelerar Hugging Face con PyTorch 2.0

Probemos la función torch.compile para acelerar los transformers de Hugging Face. Puedes hacer que tu código de Hugging Face vaya más rápido con un decorador de una sola línea.

Nota: con torch.compile() hemos observado un aumento de rendimiento del 30% al 200% en entrenamiento - TorchDynamo Performance Dashboard.

En este ejemplo aplicaremos torch.compile a un modelo de lenguaje grande "dolly-v2-3b" para acelerar la inferencia. Para ejecutar el código en Google Colab, primero instalamos las librerías necesarias de Python.

%%capture
%pip install transformers accelerate xformers

Después descargaremos y cargaremos el tokenizer y el modelo de lenguaje usando Hugging Face transformers. A continuación, añadiremos el nn.Module a la función torch.compile().

import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    GenerationConfig,
    pipeline,
)


tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-3b")
model = AutoModelForCausalLM.from_pretrained(
    "databricks/dolly-v2-3b", device_map="auto", torch_dtype=torch.bfloat16
)
model = torch.compile(model) #solo hace falta esta línea

En el último paso, convertimos el texto en tokens con el tokenizer, lo pasamos a model.generate y después decodificamos la salida generada a texto con tokenizer.batch_decode.

prompt = "I love you because"
inputs = tokenizer(prompt, return_tensors="pt").to(device="cuda:0")


# Generar
generate_ids = model.generate(inputs.input_ids, max_length=50)
tokenizer.batch_decode(
    generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
)[0]

Como vemos, "dolly-v2" ha completado la frase añadiendo: "you are a good person…"

"I love you because you are a good person.
You are kind, you help others, you are honest, you are loyal, you are humble, you are humble, you are humble.
You are a good person."

También funciona con el pipeline de Hugging Face. Solo tienes que indicar el tipo de tarea, el modelo y el tokenizer.

generator = pipeline("text-generation", model= model,tokenizer=tokenizer)
generator("What is the name of Germany's Capital?")

Salida:

[{'generated_text': "What is the name of Germany's Capital? 
The name of Germany's Capital is Berlin."}]

La función de compilación funciona de forma nativa con las librerías de Python transformers, accelerate y TIMM.

Si no tienes claro por dónde empezar tu carrera profesional en deep learning e IA, el itinerario Machine Learning Scientist with Python es un excelente punto de partida. Te ofrece una visión completa de las habilidades clave de Python que pueden ayudarte a conseguir un puesto como científico de machine learning.