PyTorch 2.0 chegou: tudo o que sabemos

PyTorch é um framework de deep learning open source, liderado pela comunidade, que oferece uma forma flexível e eficiente de construir modelos de machine learning. Ele tem uma interface amigável, amplo suporte da comunidade e integração perfeita com o ecossistema Python.

O PyTorch 2.0 trouxe mudanças fundamentais nas operações centrais do compilador, mantendo o mesmo nível de familiaridade e facilidade para desenvolvedores. Esse aprimoramento promete desempenho acelerado e suporte ampliado a Dynamic Shapes e Distributed.

O que há de novo no PyTorch 2.0?

O PyTorch está movendo partes de C++ de volta para Python, deixando tudo mais rápido e fácil de customizar. Com a versão 2.0, foi introduzido o torch.compile, que mudou a forma como o PyTorch opera no nível do compilador. Esse recurso é opcional e não afeta seu código antigo.

Compilação no PyTorch 2.0

Para dar base sólida ao recurso torch.compile, novas tecnologias foram apresentadas:

• TorchDynamo. Um compilador JIT (Just-in-Time) em nível de Python projetado especificamente para acelerar o PyTorch. Integrado à API de avaliação de frames do CPython, ele modifica dinamicamente o bytecode de Python em tempo de execução, permitindo executar o código mais rápido.
• AOTAutograd. Um toolkit que ajuda desenvolvedores a acelerar o treino de modelos no PyTorch. Ele traça antecipadamente os grafos forward e backward. Além disso, o AOTAutograd oferece mecanismos simples para compilar os grafos extraídos usando compiladores de deep learning de última geração.
• PrimTorch. Ao reduzir significativamente o número de operadores do PyTorch de mais de 2000 para um conjunto conciso de cerca de 250 operadores primitivos, o PrimTorch tornou muito mais simples desenvolver recursos ou backends para o PyTorch.
• TorchInductor. Um compilador de deep learning nativo do PyTorch que mapeia automaticamente modelos em código gerado para múltiplos aceleradores e backends. O TorchInductor usa o OpenAI Triton como bloco de construção para aceleração em GPU.

Todas essas novas tecnologias são escritas em Python e dão suporte a Dynamic Shapes. Isso faz o novo PyTorch executar código mais rápido, com mais flexibilidade e fácil de ajustar, reduzindo a barreira de entrada.

Exemplos de código

Vamos ver uma implementação rápida e simples do PyTorch Compiler.

Sem torch.compile

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

Com torch.compile

Para turbinar o desempenho do modelo, basta envolver o modelo com torch.compile e você terá um modelo compilado. É plug-and-play.

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

Aprenda a construir modelos de deep learning com a biblioteca PyTorch fazendo o curso Deep Learning with PyTorch.

Você pode treinar seu modelo sem mudar nada.

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

Ou rodar inferência do modelo.

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

A função torch.compile() vem com parâmetros adicionais:

1. mode: permite especificar em que o compilador deve focar a otimização durante a compilação.
2. dynamic: habilita o caminho de código para Dynamic Shapes.
3. fullgraph: compila o programa em um único grafo.
4. backend: por padrão usa o TorchInductor, mas você pode indicar outros backends de compilação disponíveis.

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

Aprenda o básico da API do PyTorch e crie uma rede neural simples do zero com o nosso tutorial de PyTorch.

Benchmark

Para começar com o PyTorch 2.0, os desenvolvedores usaram 163 modelos open source (46 do HuggingFace Transformers, 61 do TIMM e 56 do TorchBench) para criar benchmarks de desempenho do novo recurso de compilação. O benchmark inclui tarefas como classificação de imagens, geração de imagens, modelagem de linguagem, sistemas de recomendação e aprendizagem por reforço.

O resultado mostra um ganho de desempenho significativo durante o treino em GPUs NVIDIA A100.

Observação: no momento, o backend padrão dá suporte apenas a CPUs e às séries de GPUs Nvidia Volta e Ampere.

Benchmark do compilador PyTorch na GPU NVIDIA A100

É só o começo — nas próximas atualizações, você verá mais ganhos de performance e escalabilidade.

Como instalar o PyTorch 2.0

Você pode instalar a versão mais recente do PyTorch facilmente com pip.

Copie e cole o comando abaixo no seu terminal.

Para GPUs: CUDA 11.8

As versões mais novas de GPU têm mostrado desempenho muito 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 verificação:

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

Acelerando o Hugging Face com PyTorch 2.0

Vamos usar a função torch.compile para acelerar transformers do Hugging Face. Você pode deixar seu código do Hugging Face mais rápido com um decorator de uma linha.

Observação: com torch.compile(), já vimos um ganho de 30% a 200% no desempenho de treino — veja o Painel de Desempenho do TorchDynamo.

Neste exemplo, vamos aplicar torch.compile a um modelo de linguagem grande "dolly-v2-3b" para inferência mais rápida. Para rodar o código no Google Colab, primeiro instalamos as bibliotecas Python necessárias.

%%capture
%pip install transformers accelerate xformers

Em seguida, baixamos e carregamos o tokenizer e o modelo de linguagem usando o Hugging Face transformers. Depois disso, adicionamos o nn.Module à função 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) #only line of code is required

Na etapa final, convertemos o texto em tokens com o tokenizer, passamos para model.generate e então decodificamos a saída gerada em texto usando tokenizer.batch_decode.

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


# Generate
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 podemos ver, o "dolly-v2" completou a frase adicionando: "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."

Também funciona com a pipeline do Hugging Face. Basta informar o tipo de tarefa, o modelo e o tokenizer.

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

Saída:

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

A função de compilação funciona de primeira com as bibliotecas Python transformers, accelerate e TIMM.

Se você não sabe por onde começar sua carreira em deep learning e IA, a trilha de carreira Machine Learning Scientist with Python é um excelente ponto de partida. Ela oferece uma visão completa das principais habilidades em Python para conquistar uma vaga como cientista de machine learning.