O Guia Prático de PyTorch Avançado

Introdução

Dominar o PyTorch já não significa apenas familiaridade com as funcionalidades; trata-se de operar um fluxo de trabalho de engenharia repetível, onde o código de treino se mantém rápido, escalável e recuperável sob cargas de trabalho de produção realistas. Funcionalidades como torch.compile, torch.profiler, DDP/FSDP e Distributed Checkpointing são armas poderosas no arsenal das equipas de ML e infra para manter o treino rápido. No entanto, são inúteis se não forem aplicadas na sequência correta e validadas de forma rigorosa.

Neste artigo, partilhamos um fluxo de trabalho recomendado: linha de base → compilar → perfil → dimensionar → ponto de verificação. Iremos abordar o que medir antes de otimizar, armadilhas comuns a evitar com o compilador e o perfilador, regras de decisão a seguir para DDP vs FSDP, e como implementar pontos de verificação tolerantes a falhas para execuções de treino em múltiplos nós.

Principais Conclusões

• Aborde a otimização de desempenho do PyTorch como um processo de engenharia iterativo, em vez de uma lista de funcionalidades: aplicar iterativamente base → compilar → perfilar → escalar → criar checkpoint leva a um desempenho mais sustentável do que ativar otimizações isoladamente.
• A menos que tenha uma referência estável com uma única GPU e execução eager, com taxa de transferência conhecida e corretude verificada, não pode comparar o desempenho de forma eficaz nem depurar problemas.
• Nunca otimize antes de estabelecer uma referência correta: Uma linha de base estável com GPU única em modo eager, com rendimento medido e corretude verificada, é essencial para comparações de desempenho e depuração significativas.
• Use torch.compile deliberadamente, não cegamente: Acompanhe as interrupções do gráfico, gerencie formas, aqueça antes de medir e verifique se o desempenho em estado estável realmente melhora em relação à execução eager.
• Perfil para orientar decisões, não para confirmar suposições: Use torch.profiler para tomar decisões explícitas, resolver bloqueios de CPU, pontos críticos do kernel, re-rastreamento de formas e validar a sobrecarga de comunicação em execuções distribuídas.
• Pontos de verificação de design para falhas desde o primeiro dia: O Checkpointing Distribuído (com salvamentos assíncronos opcionais) deve capturar o estado completo do treino, suportar redistribuição entre diferentes números de GPUs e ser testado regularmente através de exercícios de restauração para garantir verdadeira tolerância a falhas.

Referência de Base: Estabelecer um Ponto de Referência

Comece por ter um exemplo de treino a funcionar com uma única GPU. Este será o seu ponto de referência tanto para a funcionalidade como para o desempenho. Defina o modelo, o dataloader e o ciclo de treino. Assegure-se de que tudo funciona de ponta a ponta em modo eager (sem compilação, um processo). Aqui está um exemplo de como o ciclo de treino pode ser:

import torch
import torch.nn as nn
# Dummy model and data for illustration
model = nn.Sequential(nn.Linear(100, 50), nn.ReLU(), nn.Linear(50, 10))
data = torch.randn(32, 100)
targets = torch.randint(0, 10, (32,))
# Baseline forward + backward
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
outputs = model(data)
loss = nn.CrossEntropyLoss()(outputs, targets)
loss.backward()
optimizer.step()

Execute algumas iterações de treino e meça o desempenho (por exemplo, amostras/segundo). Ter uma linha de base correta garante que o seu modelo converge e fornece um ponto de comparação para futuras otimizações. Neste ponto, a correção e o desempenho da linha de base são prioridades: assegure-se de que a GPU está a ser utilizada (verifique com nvidia-smi ou logs do torch), e que não existam gargalos óbvios significativos (loadgen bloqueado nos dados, operações excessivas na CPU, etc.). Não avance para a compilação e outras otimizações até ter uma linha de base estável.

Lista de Verificação Base

• Correcção funcional: O modelo treina e produz os resultados esperados numa única GPU.
• Desempenho básico registado: por exemplo, tempo por lote ou percentagem de utilização da GPU.
• Sem gargalos óbvios: O pipeline de dados mantém a GPU ocupada (sem longos períodos de inatividade na linha de base).

Uma vez que tenha implementado uma boa base, pode começar a adicionar funcionalidades avançadas de forma incremental, começando pelo compilador principal do PyTorch 2.x para acelerar o treino.

Compilar: Acelere com torch.compile no PyTorch

O PyTorch 2 permitiu aos utilizadores compilar modelos com JIT para desempenho usando torch.compile. Envolva o seu modelo ou função com torch.compile, e o PyTorch gera código otimizado nos bastidores.

Veja como é fácil abaixo:

# Switch to compiled mode
model = torch.compile(model)  # uses default backend 'inductor'

Chamar model(data) irá automaticamente usar um caminho de código otimizado. As primeiras iterações irão compilar o seu modelo em tempo real; iterações futuras irão usar os kernels otimizados diretamente. O PyTorch 2.9 guarda automaticamente os resultados da compilação em cache para melhorar execuções futuras (mesmo entre processos).

No entanto, se quiser obter o máximo desempenho das compilações, há algumas coisas que deve saber sobre interrupções de gráfico e formas dinâmicas.

Quebras de gráfico

Os cortes no gráfico acontecem sempre que o compilador não é capaz de capturar uma parte do seu código num único gráfico (como no caso de fluxos de controlo Python dependentes de dados). Por defeito, torch.compile irá automaticamente recorrer à execução imediata para as partes não suportadas (quebrando o gráfico) e compilar o restante. Isto garante que o seu código não irá falhar. No entanto, cada corte no gráfico impede que parte do código seja otimizada (essa parte corre puramente em Python, não no gráfico fundido). torch.compile(fullgraph=True) é útil para depuração durante o desenvolvimento de cortes no gráfico, pois gera um erro se qualquer parte do modelo não puder ser compilada. Por exemplo:

model = torch.compile(model, fullgraph=True)
try:
    model(data)
except Exception as e:
    print("Graph break:", e)
model = torch.compile(model, fullgraph=True)
try:
    model(data)
except Exception as e:
    print("Graph break:", e)

Isto causará um erro na primeira operação não suportada e indicar-lhe-á as partes que precisam de ser reestruturadas. A mensagem de erro normalmente inclui algumas dicas (ou um URL) sobre o que causou a falha e como a evitar. Os registos do PyTorch podem ser úteis: se executar o seu script com a variável de ambiente TORCH_LOGS=“graph_breaks”, ele imprimirá as razões e locais onde ocorrem quebras no grafo. Utilize a informação fornecida para reescrever/remover operações do lado do Python que impedem a compilação (como substituir operações Python list() ou código if/else dependente de dados por operações com tensores, ou removê-las completamente).

Formas dinâmicas

Por padrão, o torch.compile irá especializar o gráfico compilado para os formatos que encontra. Se o seu modelo estiver a correr com entradas de tamanhos variados, ele irá recompilar o gráfico cada vez que encontrar um novo tamanho. Esta recompilação tem algum overhead. No PyTorch 2.9, adicionámos uma flag dinâmica. Especificar dynamic=True antecipadamente tenta gerar um único kernel que possa lidar com múltiplos tamanhos (tamanhos simbólicos). Por exemplo:

model = torch.compile(model, dynamic=True)

Isto instrui o compilador a rastrear um gráfico generalizado, o que faz um compromisso entre recompilações à medida que os tamanhos mudam (como comprimentos de sequências variáveis). Por outro lado, passar dynamic=False nunca produzirá kernels dinâmicos e forçará a especialização para formas exatas (o que pode resultar em código mais rápido se verdadeiro). Definir dynamic=None é o padrão a partir da versão 2.9 e diz ao compilador para autodetectar se houve múltiplas recompilações e, em seguida, mudar para usar um kernel dinâmico em tempo real.

• Deixe o parâmetro dinâmico no valor predefinido ou defina-o como Falso se os formatos de entrada forem constantes ou mudarem raramente, para uma especialização máxima.
• Defina dynamic=True se estiver a receber recompilações frequentes devido a variações de forma (modelos de NLP a processar frases de diferentes comprimentos, etc.) e quiser evitar atingir o limite de recompilação (que é 8 por defeito antes de desistir e reverter para a execução imediata).

Exemplo Compilado

Vamos compilar um modelo simples e demonstrar como lidar com quebras de gráfico e formas dinâmicas:

import torch
# Sample model with a potential graph break (data-dependent control flow)
class ToyModel(torch.nn.Module):
    def forward(self, x):
        # Example: data-dependent branching (not traceable by Dynamo)
        if x.sum() > 0:
            return x * 2
        else:
            return x
model = ToyModel()
# Attempt full-graph compilation to catch breaks
try:
    model = torch.compile(model, fullgraph=True)
    output = model(torch.randn(4, 4))
except Exception as e:
    print("Graph break detected:", e)
    # Rewrite model or accept partial graph compilation
    model = torch.compile(model, fullgraph=False)  # fallback to allow breaks

No código acima, o if dependente dos dados irá gerar um erro de interrupção do gráfico ao usar fullgraph=True. Capturamos esse erro, imprimimo-lo, e depois recompilamos com o comportamento padrão de permitir interrupções no gráfico. Na prática, você corrigiria o seu modelo (removendo a bifurcação dependente dos dados neste caso) até conseguir compilar com fullgraph=True sem exceções – o que significa que todo o modelo é compilável estaticamente como um único gráfico para máxima velocidade.

Compilar Lista de Verificação

• Envolva o modelo com torch.compile(): Pequena alteração no código que frequentemente leva a melhorias significativas de velocidade. Escolha o backend “inductor” por defeito, a menos que tenha uma razão para não o fazer.
• Aquecimento: Faça algumas iterações antes de cronometrar – as primeiras execuções incluem o tempo de compilação.
• Verificar falhas no gráfico: No desenvolvimento, tente executar com fullgraph=True e defina a variável de ambiente TORCH_LOGS=“graph_breaks” para identificar partes do código que não são suportadas. Remova/refatore esses usos.
• Ajustar formas dinâmicas: Se se vir a ter muitas recompilações (procure por “recompiling” em stdout/stderr/logs), pode considerar dynamic=True para gerar um gráfico com formas flexíveis.
• Acelerar monitorização: Meça o rendimento em comparação com a sua linha de base. A primeira iteração provavelmente será mais lenta, mas a velocidade em regime estacionário deverá superar a sua linha de base inicial. Caso contrário, execute com dicas de desempenho ativadas (por exemplo, defina TORCH_LOGS="perf_hints ").

Após compilar o nosso modelo e alcançar um desempenho mais rápido, deveríamos analisar a sua execução para identificar quaisquer gargalos remanescentes e determinar áreas adicionais para otimização.

Perfil: Diagnosticar Estrangulamentos com torch.profiler

Os gargalos de desempenho podem ainda existir mesmo após a compilação - seja uma GPU subutilizada, gargalos de I/O ou kernels não otimizados. O PyTorch vem com um Profiler incorporado para recolher traços de execução e identificar esses gargalos. No PyTorch 2.9, o torch.profiler pode rastrear atividades da CPU e GPU, registar shapes e integrar-se com o Chrome Trace Viewer ou o TensorBoard para visualização. O torch.profiler também pode capturar traços de forma assíncrona. Desta forma, é possível perfilar partes do treino sem interromper o fluxo do programa.

import torch
import torch.nn as nn
import torch.optim as optim
import torch.profiler
# -----------------------------
# Device setup
# -----------------------------
device = "cuda" if torch.cuda.is_available() else "cpu"
# -----------------------------
# Model definition
# -----------------------------
model = nn.Sequential(
    nn.Linear(100, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
).to(device)
# -----------------------------
# Optimizer and loss
# -----------------------------
optimizer = optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()
# -----------------------------
# Data iterator (synthetic)
# -----------------------------
def data_generator(batch_size=32):
    while True:
        inputs = torch.randn(batch_size, 100, device=device)
        targets = torch.randint(0, 10, (batch_size,), device=device)
        yield inputs, targets
data_iter = data_generator()
# -----------------------------
# Training step
# -----------------------------
def train_step(batch):
    model.train()
    inputs, targets = batch
    optimizer.zero_grad(set_to_none=True)
    outputs = model(inputs)
    loss = loss_fn(outputs, targets)
    loss.backward()
    optimizer.step()
    return loss
# -----------------------------
# Warm-up phase
# -----------------------------
for _ in range(5):
    train_step(next(data_iter))
# -----------------------------
# Profiling phase
# -----------------------------
with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    record_shapes=True,
    profile_memory=True,
) as prof:
    for _ in range(3):
        train_step(next(data_iter))
        prof.step()
# -----------------------------
# Export trace
# -----------------------------
prof.export_chrome_trace("trace.json")
print("Profiler trace saved to trace.json")

Este script de exemplo mostra como definir um ciclo de treino PyTorch mínimo (mas completo) e usar o torch.profiler para capturar o desempenho constante da CPU e GPU. Ele constrói uma rede neural de exemplo, cria dados sintéticos com um iterador e encapsula uma iteração de treino na função train_step().

Executamos primeiro algumas iterações de aquecimento (para não perfilar a sobrecarga de inicialização/caching única). Depois, analisamos alguns passos de treino, registando as formas dos operadores e a utilização da memória. Também exportamos o rastreio para o Chrome (trace.json), que pode ser carregado no navegador Chrome para inspecionar a utilização da GPU, lançamentos de kernels, sobreposição CPU–GPU e potenciais gargalos de desempenho. Ativámos record_shapes=True e profile_memory=True para acompanhar as formas das entradas e as alocações de memória para depurar erros OOM ou ineficiências.

Dica: Abra chrome://tracing no Google Chrome e carregue o trace.json para obter uma visualização em linha de tempo.

Para deteção automática de gargalos, pode utilizar torch.utils.bottleneck ou torch.profiler.schedule para capturar instantâneos. Por exemplo, fazer profiling de alguns passos em cada época usando um agendador:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import torch.profiler
# -----------------------------
# Device
# -----------------------------
device = "cuda" if torch.cuda.is_available() else "cpu"
# -----------------------------
# Model
# -----------------------------
model = nn.Sequential(
    nn.Linear(100, 256),
    nn.ReLU(),
    nn.Linear(256, 10),
).to(device)
# -----------------------------
# Optimizer and loss
# -----------------------------
optimizer = optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()
# -----------------------------
# Dataset and DataLoader
# -----------------------------
num_samples = 10_000
batch_size = 32
inputs = torch.randn(num_samples, 100)
targets = torch.randint(0, 10, (num_samples,))
dataset = TensorDataset(inputs, targets)
data_loader = DataLoader(
    dataset,
    batch_size=batch_size,
    shuffle=True,
    num_workers=2,
)
# -----------------------------
# Training step
# -----------------------------
def train_step(batch):
    model.train()
    inputs, targets = batch
    inputs = inputs.to(device, non_blocking=True)
    targets = targets.to(device, non_blocking=True)
    optimizer.zero_grad(set_to_none=True)
    outputs = model(inputs)
    loss = loss_fn(outputs, targets)
    loss.backward()
    optimizer.step()
    return loss
# -----------------------------
# Warm-up (outside profiler)
# -----------------------------
for i, batch in enumerate(data_loader):
    train_step(batch)
    if i >= 5:
        break
# -----------------------------
# Profiling with schedule
# -----------------------------
with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    schedule=torch.profiler.schedule(
        wait=1,
        warmup=1,
        active=3,
        repeat=2,
    ),
    on_trace_ready=torch.profiler.tensorboard_trace_handler("./prof_log"),
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    for step, batch in enumerate(data_loader):
        loss = train_step(batch)
        prof.step()

        if step >= 50:
            break

Isto irá iniciar de forma assíncrona dois períodos de rastreio (cada um com 3 passos) após ignorar um período de espera/aquecimento. O callback on_trace_ready depois envia os rastreios que pode visualizar na aba de perfil do TensorBoard sem parar o treino. A captura de rastreio assíncrona permite-lhe perfilar de forma intermitente (por exemplo, a cada 3 passos numa janela de 50 passos) para reduzir a sobrecarga, enquanto ainda captura dados de desempenho.

Após a análise de perfil, interprete os resultados para tomar medidas:

• Limitado pela CPU (muitas pausas enquanto a GPU espera): Considere se consegue transferir mais trabalho para a GPU, ou sobrepor o carregamento de dados com a computação usando carregamento assíncrono de dados (por exemplo, passando `workers` para o DataLoader ou pré-processamento na GPU). Além disso, verifique se não há laços Python a executar trabalho durante o treino; se for o caso, tente removê-los ou `compilá-los`.
• Os kernels da GPU são o gargalo: Pode otimizar essas operações (operações fundidas, menor precisão, etc.). O guia de desempenho do PyTorch recomenda manter operações de alto nível (como chamadas torch.nn.functional) para permitir que as bibliotecas as otimizem. No entanto, se notar uma operação em particular que seja muitas vezes mais lenta, verifique se isso é esperado (carregar FFTs grandes pode ser lento) ou se existem formas alternativas de a calcular.
• Formas múltiplas a causar retracing: considere agrupar os inputs por tamanho, ou use dynamic=True como mencionado acima.
• Assegure-se de que não existem gargalos externos: Por exemplo, se estiver a utilizar treino distribuído, verifique se a comunicação (all-reduce) não é o gargalo - isso seria indicado pelos kernels NCCL ou pela CPU à espera da rede.

Lista de Verificação de Perfil

• Aquecimento antes do profiling: Evite medir o overhead de compilação ou inicialização preguiçosa.
• Capture tanto a CPU como a GPU: Capture tanto a CPU como a GPU (actividades=[CPU, CUDA]) para compreender a sua interação.
• Gravar formas e memória: Ajuda a identificar problemas relacionados com formas e picos de utilização de memória.
• Utilize agendamento para execuções longas: Utilize profiler.schedule para capturar janelas de tempo curtas periodicamente. Isto reduz a sobrecarga do profiler e ainda assim proporciona visibilidade sobre o comportamento do programa.
• Analise a utilização: Certifique-se de que as suas GPUs estão a ser totalmente utilizadas. Se não estiverem, verifique se o limitador é a CPU ou o I/O.
• Identificar operações principais: O profiler pode gerar uma lista de operações ordenadas por métricas importantes, como tempo próprio. Determine quais operações estão a levar mais tempo e otimize-as (melhores algoritmos, aumentar o tamanho do lote se as GPUs estiverem subutilizadas).

Com a aplicação otimizada numa única máquina, o próximo desafio é expandir para tirar partido de múltiplas GPUs ou nós, o que introduz os paradigmas de treino distribuído do PyTorch.

Escala: Treino Distribuído via DDP ou FSDP

Quando a sua carga de trabalho ou modelo crescer para além de uma única GPU, será necessário expandir o treino para várias GPUs. O PyTorch oferece dois métodos principais para treinar modelos em múltiplas GPUs: Data Parallel Distribuído e Data Parallel Totalmente Fracionado. Embora ambos os métodos utilizem paralelismo de dados internamente (cada processo treina um subconjunto diferente dos seus dados), eles têm abordagens diferentes para gerir os parâmetros do modelo na memória. Vamos explorar quando escolher cada método e como configurá-los tanto para treino em nó único como em múltiplos nós com o torchrun.

Paralelismo de Dados Distribuído

O DDP mantém uma cópia de todo o modelo em cada GPU e sincroniza as cópias através da redução total dos gradientes após cada passo. Ideal para modelos que cabem facilmente na memória de uma única GPU, o DDP é conceitualmente simples: inicializar um grupo de processos e depois envolver o modelo com torch.nn.parallel.DistributedDataParallel.

import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, TensorDataset
from torch.utils.data.distributed import DistributedSampler
# -----------------------------
# Setup distributed (safe)
# -----------------------------
def setup_distributed():
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    if world_size > 1:
        dist.init_process_group(backend="nccl")
        local_rank = int(os.environ["LOCAL_RANK"])
        torch.cuda.set_device(local_rank)
        device = torch.device(f"cuda:{local_rank}")
        distributed = True
    else:
        local_rank = 0
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        distributed = False

    return distributed, local_rank, device

def cleanup_distributed(distributed):
    if distributed:
        dist.destroy_process_group()

# -----------------------------
# Main training
# -----------------------------
def main():
    distributed, local_rank, device = setup_distributed()

    # -----------------------------
    # Model
    # -----------------------------
    model = nn.Sequential(
        nn.Linear(100, 256),
        nn.ReLU(),
        nn.Linear(256, 10),
    ).to(device)
    if distributed:
        model = DDP(model, device_ids=[local_rank])
    # -----------------------------
    # Optimizer and loss
    # -----------------------------
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    loss_fn = nn.CrossEntropyLoss()
    # -----------------------------
    # Dataset
    # -----------------------------
    num_samples = 20_000
    batch_size = 32
    inputs = torch.randn(num_samples, 100)
    targets = torch.randint(0, 10, (num_samples,))
    dataset = TensorDataset(inputs, targets)
    if distributed:
        sampler = DistributedSampler(dataset, shuffle=True)
    else:
        sampler = None
    train_loader = DataLoader(
        dataset,
        batch_size=batch_size,
        sampler=sampler,
        shuffle=(sampler is None),
        num_workers=2,
        pin_memory=True,
    )

    # -----------------------------
    # Training loop
    # -----------------------------
    epochs = 3
    for epoch in range(epochs):
        if distributed:
            sampler.set_epoch(epoch)

        for step, (x, y) in enumerate(train_loader):
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)

            optimizer.zero_grad(set_to_none=True)
            outputs = model(x)
            loss = loss_fn(outputs, y)
            loss.backward()
            optimizer.step()
            if step % 50 == 0 and local_rank == 0:
                print(
                    f"[Epoch {epoch}] Step {step} | "
                    f"Loss {loss.item():.4f}"
                )
    cleanup_distributed(distributed)
# -----------------------------
# Entry point
# -----------------------------
if __name__ == "__main__":
    main()

Lançar um trabalho DDP é feito com o torchrun. Por exemplo, se quiser lançar um trabalho num único nó com 4 GPUs:

torchrun --nproc_per_node=4 train.py

O comando irá criar 4 processos (um para cada GPU) e iniciar as variáveis de ambiente apropriadas (rank, dimensão do mundo, etc.). Especificar múltiplos nós requer que passe --nnodes e os detalhes da rede:

# On node 0:
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=0 --master_addr="<IP of node0>" --master_port=12345 train.py
# On node 1:
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=1 --master_addr="<IP of node0>" --master_port=12345 train.py

Paralelismo de Dados Totalmente Fragmentado (FSDP)

O FSDP vai um passo além ao dividir os parâmetros do modelo e os estados do otimizador entre as GPUs, em vez de os replicar completamente em cada GPU. Utilizar o FSDP no PyTorch 2.9 é simples com o invólucro FullyShardedDataParallel.

Pontos principais

• Envolva o modelo (ou submódulos) com FSDP antes da criação do otimizador.
• Use torch.cuda.set_device(rank) e mova o modelo para a GPU como com DDP.
• Opcional: forneça uma política de auto-embalamento se desejar dividir por granularidade de módulo (por exemplo, embalar cada bloco transformador separadamente). Por defeito, FSDP(model) irá dividir o modelo inteiro como uma única unidade de shard (tecnicamente uma unidade FSDP), o que o torna subótimo para modelos muito profundos.

Isto irá fragmentar recursivamente quaisquer submódulos com mais de 100k parâmetros em os seus próprios fragmentos FSDP. Se tiver poucos fragmentos (por exemplo, se todo o modelo for um único fragmento), não irá obter poupança de memória entre camadas. No entanto, muitos fragmentos aumentam o overhead de comunicação. Vamos resumir as diferenças entre DDP e FSDP:

Para fins de exemplo, pode começar com DDP. Se encontrar problemas de memória ou precisar de modelos maiores, considere usar FSDP. Note que o FSDP no PyTorch 2.9 evoluiu; por defeito, ele “fragmenta tudo” (ZeRO Stage 3) e utiliza outros padrões, como limit_all_gathers=True, para evitar picos de memória inesperados.

Lista de Verificação para Escalamento

• Garantir determinismo: Certifique-se de que está a definir a mesma semente aleatória em todos os ranks (torch.manual_seed(seed for ran offset) se precisar de reprodutibilidade.
• Usar DistributedSampler: Divida os seus dados em múltiplos ranks e chame set_epoch() em cada época para que sejam embaralhados de forma diferente a cada época.
• DDP: Os gradientes são automaticamente somados entre as GPUs após chamar backward() nos seus modelos envoltos em DDP. Não se esqueça de aumentar a taxa de aprendizagem ou o tamanho do batch ao aumentar o número de GPUs!
• FSDP: Envolva os seus módulos antes de criar o otimizador. Experimente com a auto_wrap_policy para evitar um único fragmento enorme (provável se tiver um modelo muito profundo). Mantenha atenção à GPU – assegure-se de que o tamanho do fragmento realmente cabe.
• Backend de comunicação: Para múltiplos nós, certifique-se de que tem uma rede configurada (normalmente NCCL). Defina a variável de ambiente NCCL_P2P_LEVEL=NVL se tiver NVLink, etc., para otimizar o NCCL.
• Acumulação de gradientes: Se a utilizar, assegure-se de que, com o FSDP, chama model.no_sync() se estiver a saltar a sincronização de gradientes em determinadas iterações (de forma semelhante ao uso do DDP).

Ponto de Verificação: Recuperar o Treino de Forma Fiável com Pontos de Verificação Distribuídos

Em sessões de treino longas ou de grande escala, é importante guardar pontos de verificação. Pode querer retomar o treino após uma falha, ajustar a formação a partir do meio do treino ou inspecionar resultados intermédios. O PyTorch 2.9 inclui primitivas convenientes para a criação distribuída de pontos de verificação (DCP). Isto permite aos utilizadores guardar/carregar de forma mais eficiente e robusta do que as abordagens tradicionais com torch.save. Vamos comparar ambas as abordagens antes de mergulhar na abordagem recomendada.

torch.save/torch.load tradicionais

O método tradicional num script de treino com uma única GPU ou mesmo com DDP irá guardar o state_dict() do modelo num ficheiro (normalmente apenas no rank 0, já que todos os processos terão acesso a este ficheiro). Por exemplo:

# On rank 0 only:
torch.save(model.state_dict(), "checkpoint.pt")
and to load:
model.load_state_dict(torch.load("checkpoint.pt"))

Embora isto seja aceitável para modelos/configurações pequenos, onde um único processo pode manter todo o estado do modelo na memória (normalmente o rank 0), existem problemas ao usar esta abordagem para treino distribuído, onde o modelo pode estar fragmentado entre processos.

Ponto de Verificação Distribuído (DCP): o módulo torch.distributed.checkpoint no PyTorch (adicionado desde a versão 2.1+, amadurecido até à 2.9) aborda estas questões. Ele paraleliza a gravação: cada rank escreve a sua parte do estado do modelo, produzindo múltiplos ficheiros (pelo menos um por rank) que constituem o ponto de verificação. Também realiza o re-sharding ao carregar: é possível guardar em N ranks e carregar em M ranks, e o DCP irá automaticamente agregar/dividir conforme apropriado.

O padrão típico de utilização envolve envolver o seu modelo (e o otimizador, se necessário) numa estrutura com estado que fornece um state_dict. Pode fazer isto utilizando as funções utilitárias fornecidas get_state_dict e set_state_dict, que são compatíveis com FSDP:

import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
from torch.distributed.checkpoint.stateful import Stateful

# Define a stateful container for model & optimizer
class AppState(Stateful):
    def __init__(self, model, optimizer=None):
        self.model = model
        self.optimizer = optimizer
    def state_dict(self):
        model_sd, opt_sd = get_state_dict(self.model, self.optimizer)
        return {"model": model_sd, "optim": opt_sd}
    def load_state_dict(self, state):
        set_state_dict(self.model, self.optimizer, state["model"], state.get("optim"))

AppState aqui implementa a interface Stateful, por isso o DCP sabe como obter e definir o estado. Note que get_state_dict usa automaticamente o dicionário de estado distribuído do FSDP se o modelo estiver distribuído (cada rank obtém apenas a sua parte do modelo).

Poupar com a DCP

Todos os escalões executam:

app_state = AppState(model, optimizer)
dcp.save(app_state, checkpoint_id="mycheckpoint")

Isto irá gerar diretórios contendo ficheiros com o prefixo "mycheckpoint" (mycheckpoint/mpirank_0_0.pdparams). Cada rank grava o seu fragmento em paralelo. Isto pode reduzir significativamente o tempo em comparação com um único rank a fazê-lo.

A carregar com DCP

Para retomar a partir de um ponto de verificação, precisará de um modelo com a mesma topologia (potencialmente com um número diferente de ranks). Construa o seu modelo (e o otimizador), envolva-o em AppState, e depois:

app_state = AppState(model, optimizer)
dcp.load(app_state, checkpoint_id="mycheckpoint");

Desde que o número de ranks seja diferente do que tinha quando guardou, o DCP encarregar-se-á de reorganizar os shards internamente (por exemplo, passar de 8 para 4 GPUs significa que cada rank irá carregar dois shards guardados, ou, se aumentar o número de ranks, será distribuído de forma adequada). Esta etapa de resharding no momento da carga elimina a necessidade de qualquer conversão manual de checkpoints.

Checkpointing assíncrono

Uma das funcionalidades mais interessantes do PyTorch 2.9 é salvamento asincrónico de checkpoints, que permite sobrepor a escrita do checkpoint com o treino. Utilizando dcp.async_save, a API retorna um futuro e realiza a operação intensiva de E/S em threads de fundo. A ideia básica:

save_future = dcp.async_save(app_state, checkpoint_id="chkpt_epoch10")
# ... training can continue immediately ...
save_future.wait()  # later, wait for completion (or periodically check)

async_save primeiro prepara os dados localmente na CPU (copia o estado do modelo/otimizador para buffers de memória fixos), e depois grava esses buffers no disco/rede de forma assíncrona. As GPUs podem retomar o trabalho de computação (ou seja, o treino) enquanto o ponto de verificação está a ser descarregado.

A gravação assíncrona requer temporariamente o dobro da memória para armazenar os estados do seu modelo. Cada nível aloca um buffer na CPU aproximadamente igual em tamanho à sua parte do checkpoint. Portanto, guardar modelos muito grandes de forma assíncrona através de vários níveis pode aumentar consideravelmente o uso de memória (RAM da CPU + memória fixa ~ tamanho_do_checkpoint_por_nível * número_de_níveis, no pior cenário).

Agora, vamos compilar uma comparação dos métodos de checkpointing com orientações sobre quando usar cada um:

Lista de Verificação de Checkpoint

• Inclua todo o estado necessário: Certifique-se de guardar os pesos do modelo e o estado do otimizador (e do agendador, estado do RNG, etc., se necessário) para poder retomar completamente o treino mais tarde. Se estiver a usar o método Stateful do DCP, pode facilmente passar o otimizador através de get_state_dict(…, model, optimizer).
• Teste de restauração: Assim que tiver guardado o ponto de verificação, carregue-o de volta (preferencialmente no mesmo número e em diferentes números de GPUs se distribuído) num ambiente limpo para garantir que tudo se restaura como esperado. Isto ajudá-lo-á a encontrar erros, como esquecer de guardar alguma parte do estado de treino no início.
• Gerir armazenamento: Os pontos de verificação podem crescer muito rapidamente (especialmente com pontos de verificação distribuídos em fragmentos). Implemente alguma estratégia de limpeza (manter os últimos k pontos de verificação, etc.) para evitar ficar sem espaço no disco.
• Uso asincrónico: async_save é uma funcionalidade útil, mas deve ter cuidado para não ter múltiplas operações de gravação assíncrona em curso ao mesmo tempo (a menos que tenha muito cuidado com a forma como utiliza a memória e permite sobreposições). Eventualmente, deve chamar .wait() sobre o futuro resultante para levantar quaisquer exceções ou confirmar o sucesso.
• Consistência: Ao retomar a partir de um ponto de verificação, certifique-se de que todas as classificações chamem o mesmo dcp.load. Isto é especialmente importante para FSDP, pois é necessário sincronizar o carregamento interno. Além disso, certifique-se de que o modelo esteja nos dispositivos esperados antes de carregar, se estiver a utilizar descarregamento para CPU ou outros mecanismos.

Conclusão

Ao seguir este fluxo de trabalho, irá desenvolver código de treino em PyTorch que tem bom desempenho e é robusto a falhas:

• Comece a partir de uma base funcional.
• Compilar para a máxima velocidade.
• Perfil para identificar oportunidades de otimização.
• Escale com uma estratégia adequada de paralelismo de dados ou de modelos.
• Ponto de verificação eficaz para permitir a recuperação de falhas.

Existem guias disponíveis para cada etapa. Utilize-os de forma iterativa (por exemplo, faça novamente o perfil após a etapa 4 alterar as suas características de desempenho, ou reveja as configurações de compilação na etapa 2 para ajustar ao seu novo ambiente). Como o PyTorch tem evoluído rapidamente, existem capacidades que ajudam em cada etapa. Saiba quando e como usar cada uma, e estará preparado para treinar modelos grandes com rapidez e escala, com boas estratégias de recuperação.

Perguntas Frequentes

Por que é tão importante estabelecer uma linha de base antes da otimização?

É necessário ter uma linha de base correta numa única GPU usando execução eager antes de otimizar, para ter algo em que confie quanto à corretude e desempenho. Caso contrário, as otimizações podem ocultar erros que provocam resultados incorretos. Só com uma linha de base correta pode atribuir melhorias de velocidade ou regressões a alterações específicas no código, em vez de erros enigmáticos ou ruído de medição.

Quando devo usar torch.compile e em que devo ter atenção?

torch.compile deve ser usado depois de ter uma linha de base sólida para acelerar a execução em estado estacionário. Monitore a interrupção do gráfico, aqueça antes de medir e assegure-se de que os formatos dinâmicos são tratados deliberadamente para evitar recompilações excessivas.

Como é que o torch.profiler ajuda para além de confirmar suposições?

torch.profiler revela os verdadeiros pontos críticos do seu código — seja falta de CPU, núcleos de GPU ineficientes, retracing de formas ou sobrecarga de comunicação. Isto permite concentrar os esforços de otimização onde realmente importa, com base em evidências e não em intuição.

Como escolher entre DDP e FSDP para treino distribuído?

Se o seu modelo se adapta confortavelmente a cada GPU que pretende usar, opte pelo DDP pela sua simplicidade e velocidade relativa. Se estiver limitado pela memória da GPU ou a trabalhar com modelos muito grandes, o FSDP proporcionará uma maior escalabilidade à custa de complexidade adicional.

Por que o checkpointing distribuído é preferido em relação a torch.save em execuções de grande escala?

O checkpointing distribuído paraleliza as gravações entre os ranks, suporta reshaping ao carregar e permite checkpointing assíncrono, tornando a recuperação fiável e eficiente para trabalhos de treino grandes, multi-GPU ou multi-nó.

Referências

• Guia definitivo da biblioteca PyTorch em Python
• Escalonamento do Treino de Modelos em Várias GPUs: Estratégias Eficientes com PyTorch DDP e FSDP
• Paralelismo de Dados Totalmente Fragmentado
• Ponto de Verificação Distribuído do Torch
• Tutorial do torch.compile

Obrigado por aprender com a comunidade Linux-Console.net.