Profiling no PyTorch (Parte 1): dominando o torch.profiler para otimizar seus modelos

O torch.profiler é a API oficial de profiling do PyTorch, lançada em 25 de março de 2021 com o PyTorch 1.8.1 e desenvolvida em parceria entre Meta (Facebook) e Microsoft. Ele substituiu definitivamente a antiga torch.autograd.profiler, que agora está em processo de descontinuação, com novas APIs de análise de memória como torch.cuda.memory._record_memory_history e _export_memory_snapshot sendo adotadas desde novembro de 2025. O profiler opera via gerenciador de contexto (with torch.profiler.profile(...)) e coleta dados granulares de CPU (operadores PyTorch, TorchScript, rótulos personalizados), GPU (kernels CUDA e XPU), transferências de memória (ex.: cudaMemcpy) e até tarefas assíncronas (torch.jit._fork). Recursos avançados como profile_memory=True, record_shapes=True, with_stack=True e with_flops=True permitem diagnósticos precisos de gargalos, vazamentos de memória, ineficiências por forma de tensor e estimativa de FLOPs, mas cada um adiciona sobrecarga mensurável, exigindo uso estratégico.

A arquitetura subjacente é baseada na biblioteca Kineto, que fornece rastreamento de linha do tempo de baixa sobrecarga para CPU+GPU, integrando-se ao CUPTI da NVIDIA, ao ROCm da AMD e ao Intel XPU. Os resultados podem ser exportados para JSON do Chrome (compatível com chrome://tracing e Perfetto) ou visualizados diretamente no TensorBoard via torch.profiler.tensorboard_trace_handler, com suporte nativo do plugin torch-tb-profiler. Ferramentas complementares como a Holistic Trace Analysis (HTA) processam rastreamentos Kineto para decompor o tempo da GPU em computação, comunicação, movimentação de memória e ociosidade, essencial para otimização em treinamento distribuído com DeepSpeed ou PyTorch Lightning.

O torch.profiler é crítico porque, sem ele, desenvolvedores operam no escuro: modelos podem estar subutilizando GPUs (ex.: menos de 30% de utilização), sofrendo stalls por transferências CPU-GPU ou consumindo memória excessiva sem explicação. Estudos de caso reais mostram que até 40% do tempo de treinamento em modelos como ResNet-50 e BERT pode ser perdido em overheads não identificados, como chamadas repetidas de torch.cat com formas variáveis ou kernels CUDA não fundidos. A ferramenta permite validar se otimizações como torch.compile (disponível desde PyTorch 2.0) estão efetivamente fundindo operadores (ex.: matmul + add → aten::addmm) e se o modelo está aproveitando bem as novas rodas PyPI do PyTorch 2.11 (com CUDA 13.0 embutida), que deixaram de suportar GPUs Maxwell e Pascal. Em ambientes produtivos, o uso inadequado do profiler, como ativar with_stack=True em loops longos, pode dobrar o tempo de execução, tornando o conhecimento das configurações ideais uma habilidade técnica obrigatória.

Para desenvolvedores de ML, dominar o torch.profiler significa reduzir ciclos de iteração em até 70% ao identificar rapidamente o 'culpado' real de lentidão: seja um operador aten::copy_ escondido em um DataLoader, um kernel CUDA não otimizado em uma camada de convolução customizada ou um vazamento de memória causado por referências mantidas com record_shapes=True. A API exige compreensão prática de trade-offs: schedule(wait=1, warmup=1, active=3, repeat=2) é padrão recomendado para capturar estado estacionário em loops de treinamento, evitando ruído inicial e arquivos gigantescos. Integrações com frameworks como PyTorch Lightning e DeepSpeed automatizam perfis de etapas-chave (ex.: training_step, backward), mas exigem ajuste fino, como desativar profile_memory em fases de warmup para evitar impacto na convergência. Além disso, o suporte a ROCm (AMD) e XPU (Intel) torna o profiler indispensável em ambientes heterogêneos, onde benchmarks baseados apenas em CUDA são insuficientes.