Internals do PyTorch (2019)

• Uma explicação da estrutura interna do PyTorch, servindo como guia para quem deseja contribuir com o código-base em C++ do PyTorch
• O objetivo deste texto é ajudar a entender a estrutura da biblioteca de tensores do PyTorch e a técnica de diferenciação automática (autograd), além de facilitar a navegação pelo código-base

Estrutura básica dos tensores no PyTorch

• No PyTorch, o tensor é a estrutura de dados mais fundamental
• Tensores são estruturas de dados n-dimensionais, capazes de armazenar valores escalares como ponto flutuante (float), inteiro (int) etc.
• Os tensores incluem os seguintes metadados:
  • tamanho (size): informação sobre as dimensões do tensor
  • dtype: tipo de dado armazenado (ex.: float32, int64 etc.)
  • device: local onde os dados são armazenados (CPU, CUDA etc.)
  • stride: informação de deslocamento dos dados na memória física

• Papel do stride

  • stride é usado para converter índices lógicos em posições na memória física
  • stride define deslocamentos para cada dimensão, e a posição na memória física é determinada multiplicando o índice pelos valores de stride
  • Com stride, é possível visualizar os mesmos dados de outra forma, como uma view, sem criar um novo tensor

Conceito de tensor e armazenamento (Storage)

• No PyTorch, o tensor não armazena diretamente os dados reais → os dados são gerenciados pelo armazenamento (Storage)
• Tensor = tamanho + dtype + device + stride + offset
• Vários tensores podem compartilhar um único armazenamento → suporte ao conceito de view (View)
• A separação entre armazenamento e tensor permite uso mais eficiente da memória

Processo de dispatch das operações de tensor

• No PyTorch, as operações passam por duas etapas de dispatch:
  1. Dispatch com base no tipo de device e no layout
     • Dependendo de ser um tensor de CPU ou um tensor CUDA, é executado um código de implementação diferente
  2. Dispatch com base em dtype
     • Kernels diferentes são chamados conforme o tipo de dado, como float vs. int

Modelo de extensão de tensores no PyTorch

• Três principais elementos de extensão do tensor:

  • Device: define como a memória é alocada em CPU, GPU, TPU etc.
  • Layout: define como o tensor é armazenado na memória (ex.: armazenamento contíguo, armazenamento esparso (sparse) etc.)
  • dtype: define o tipo de dado armazenado em cada elemento do tensor

• Opções de extensão:

  • É possível estender tensores modificando diretamente o código do PyTorch
  • Também é possível escrever uma classe wrapper que encapsula tensores existentes
  • Se for necessário um wrapper durante a diferenciação automática, será preciso fazer a extensão diretamente

Como funciona a diferenciação automática (Autograd)

• O PyTorch realiza diferenciação automática com base em retropropagação (reverse-mode differentiation)
• Durante a operação de forward, o grafo é criado → durante o backward, o grafo é percorrido para calcular as derivadas
• O Autograd gerencia as seguintes informações adicionais:
  • AutogradMeta: metadados conectados ao tensor e usados na retropropagação
  • Registra os resultados das operações e executa a diferenciação durante o backward

Estrutura do código do PyTorch e localização dos arquivos

• Diretórios principais no código-base do PyTorch:
  • torch/ → módulo Python (código Python)
  • torch/csrc/ → código de binding entre Python e C++, engine de diferenciação automática, compilador JIT etc.
  • aten/ → definições de operações de tensor (inclui a maior parte das operações centrais)
  • c10/ → definições de estruturas centrais, como tensor e armazenamento

Processo de execução de operações no PyTorch

• Exemplo: processo executado ao chamar torch.add():
  1. Conversão dos argumentos de Python para código C++
  2. Execução do dispatch em VariableType
  3. Execução do dispatch com base em device/layout
  4. Execução do kernel final

Processo e ferramentas para escrever kernels

• No PyTorch, um kernel é escrito seguindo etapas como:
  1. Escrita dos metadados da operação: definição da assinatura da função, devices suportados e tipos de dados
  2. Validação de entrada: verificação de dimensões, tipos etc.
  3. Alocação do tensor de saída
  4. Dispatch por dtype: execução do kernel conforme o tipo de dado
  5. Processamento paralelo: no CPU usa OpenMP, e no CUDA usa paralelização embutida
  6. Acesso aos dados e cálculo: uso de TensorAccessor, TensorIterator etc.

Principais macros de dispatch

• AT_DISPATCH_ALL_TYPES → realiza dispatch conforme o dtype
• Há suporte por macro para vários tipos de dados → permite otimização de desempenho

Dicas para otimização de desempenho e aumento da eficiência no trabalho

• Minimizar alterações em arquivos de cabeçalho → alterações nesses arquivos causam rebuild de todo o código
• Configurar um ambiente de desenvolvimento local → minimiza perda de tempo ao usar CI
• Usar ccache → pode economizar tempo de recompilação
• Usar um servidor potente → pode reduzir o tempo de compilação em C++ e build de CUDA

Guia de contribuição para o PyTorch

• Bons pontos de partida para contribuir:
  • Issues com a label triaged → issues já verificadas por desenvolvedores do PyTorch
  • Melhorias na documentação e ajuda para reproduzir bugs
  • Opiniões sobre RFCs (propostas de funcionalidade) do PyTorch
• O PyTorch cresceu graças a contribuidores de código aberto, e a participação da comunidade é bem-vinda