Internals de PyTorch (2019)

• Una explicación de la estructura interna de PyTorch, como guía para quienes quieran contribuir al codebase en C++ de PyTorch
• El objetivo de este artículo es ayudar a entender la estructura de la biblioteca de tensores de PyTorch y la técnica de diferenciación automática (autograd), además de orientarse dentro del codebase

Estructura básica de los tensores de PyTorch

• En PyTorch, el tensor es la estructura de datos más básica
• Un tensor es una estructura de datos de n dimensiones que puede almacenar valores escalares como punto flotante (float), enteros (int), etc.
• Los tensores incluyen los siguientes metadatos:
  • tamaño (size): información sobre las dimensiones del tensor
  • dtype: tipo de dato almacenado (por ejemplo, float32, int64, etc.)
  • device: ubicación donde se almacenan los datos (CPU, CUDA, etc.)
  • stride: información de desplazamiento en la memoria física de los datos

• El papel de stride

  • stride se usa para convertir índices lógicos en ubicaciones de memoria física
  • stride establece un desplazamiento para cada dimensión y determina la ubicación de memoria física multiplicando el índice por el valor de stride
  • Gracias a stride, es posible ver los mismos datos de otra forma como una view sin crear un tensor nuevo

Concepto de tensor y almacenamiento (Storage)

• En PyTorch, el tensor no guarda directamente los datos reales → los datos se administran en el almacenamiento (Storage)
• Tensor = tamaño + dtype + device + stride + offset
• Varios tensores pueden compartir un mismo almacenamiento → esto habilita el concepto de vista (View)
• La separación entre almacenamiento y tensor permite usar la memoria de forma eficiente

Proceso de dispatch en las operaciones de tensor

• En PyTorch, las operaciones pasan por dos etapas de dispatch:
  1. Dispatch basado en tipo de dispositivo y layout
     • Se ejecuta código de implementación distinto según sea un tensor de CPU o de CUDA
  2. Dispatch basado en dtype
     • Se llama a distintos kernels según el tipo de dato, como float vs. int

Modelo de extensión de tensores en PyTorch

• Los tres elementos principales de extensión de un tensor:

  • Device: define cómo se asigna la memoria en CPU, GPU, TPU, etc.
  • Layout: define cómo se almacena el tensor en memoria (por ejemplo, almacenamiento contiguo, almacenamiento disperso (sparse), etc.)
  • dtype: define el tipo de dato que se almacenará en cada elemento del tensor

• Opciones de extensión:

  • Es posible extender tensores modificando directamente el código de PyTorch
  • Es posible escribir una clase wrapper que envuelva tensores existentes
  • Si se necesita el wrapper durante la diferenciación automática, hace falta una extensión directa

Cómo funciona la diferenciación automática (Autograd)

• PyTorch realiza diferenciación automática basada en reverse-mode differentiation
• Durante la operación forward se crea un grafo → durante backpropagation se recorre ese grafo para realizar la diferenciación
• Autograd administra información adicional como la siguiente:
  • AutogradMeta: metadatos conectados al tensor que se usan en backpropagation
  • Registra los resultados de las operaciones y realiza la diferenciación durante backpropagation

Estructura del código de PyTorch y ubicación de archivos

• Directorios principales dentro del codebase de PyTorch:
  • torch/ → módulo de Python (código Python)
  • torch/csrc/ → bindings entre Python y C++, motor de diferenciación automática, compilador JIT, etc.
  • aten/ → definición de operaciones de tensor (incluye la mayoría de las operaciones principales)
  • c10/ → definición de estructuras de datos base como tensor y almacenamiento

Proceso de ejecución de operaciones en PyTorch

• Ejemplo: proceso de ejecución al llamar torch.add():
  1. Conversión de argumentos desde Python a código C++
  2. Se realiza el dispatch en VariableType
  3. Se realiza el dispatch basado en device/layout
  4. Ejecución final del kernel

Proceso y herramientas para escribir kernels

• En PyTorch, un kernel se escribe siguiendo pasos como estos:
  1. Escritura de metadatos de la operación: definir la firma de la función, dispositivos compatibles y tipos de dato soportados
  2. Validación de entrada: verificar entradas como dimensiones, tipos, etc.
  3. Asignación del tensor de salida
  4. Dispatch por dtype: ejecutar el kernel según el tipo de dato
  5. Procesamiento en paralelo: en CPU se usa OpenMP, y en CUDA se usa paralelización integrada
  6. Acceso a datos y cálculo: uso de TensorAccessor, TensorIterator, etc.

Macros principales de dispatch

• AT_DISPATCH_ALL_TYPES → realiza dispatch según el dtype
• Hay soporte mediante macros para diversos tipos de dato → esto permite optimizar el rendimiento

Consejos para optimización de rendimiento y eficiencia de trabajo

• Minimizar cambios en archivos de encabezado → al modificarlos se fuerza un rebuild completo del código
• Configurar un entorno de desarrollo local → minimiza el tiempo perdido al depender de CI
• Usar ccache → puede ahorrar tiempo de recompilación
• Usar un servidor potente → puede reducir el tiempo de compilación en C++ y de build de CUDA

Guía para contribuir a PyTorch

• Buenos puntos de inicio para contribuir:
  • Issues con la etiqueta triaged → issues ya revisados por desarrolladores de PyTorch
  • Mejoras de documentación y ayuda para reproducir bugs
  • Dar opinión sobre los RFC de PyTorch (propuestas de funcionalidad)
• PyTorch ha crecido gracias a contribuyentes de código abierto y da la bienvenida a la participación de la comunidad