Qué realmente impulsa una red neuronal: los compiladores de IA detrás de Python

Cualquiera con experiencia en compiladores y conocimientos básicos de aprendizaje automático podría nunca haber relacionado estos dos campos. Los analizadores léxicos, los analizadores sintácticos, los árboles de expresiones y el LLVM IR pertenecen a un mundo, mientras que las redes neuronales pertenecen a otro. Se encuentran en los compiladores de IA/ML, y la primera sorpresa es esta: un modelo de aprendizaje profundo en realidad no funciona con Python.

Eso suena extraño, porque una línea como la siguiente parece indicar que Python se encarga de todo:

output = model(input)

Pero Python no realiza por sí mismo miles de millones de multiplicaciones de matrices. Debajo de esto ocurre un proceso mucho más largo.

De una fórmula a instrucciones de hardware

Imagínese una pequeña red que va desde la entrada, pasa por una multiplicación de matrices y continúa:

Input
  ↓
Matrix Multiplication
  ↓
ReLU
  ↓
Matrix Multiplication
  ↓
Output

En Python, su núcleo podría verse así:

y = relu(x @ W + b)

Parece simple, pero piense en lo que necesita la máquina. Una CPU o GPU no tiene concepto de esta función:

relu()

o de este operador:

x @ W

tal como aparecen en el código. El hardware ejecuta las instrucciones. En algún momento, la descripción de alto nivel del modelo debe transformarse en algo que el chip pueda ejecutar de manera eficiente, y esa transformación es la función de los compiladores de IA/ML.

Una imagen simplificada de la cadena, desde el modelo hasta el procesador:

AI Model
   ↓
PyTorch / TensorFlow
   ↓
Computational Graph
   ↓
Intermediate Representation (IR)
   ↓
Optimisations
   ↓
Lowering
   ↓
Hardware-specific Code
   ↓
Machine Instructions
   ↓
CPU / GPU / NPU

Hay muchas etapas, así que vale la pena abordarlas una por una.

Qué hace el framework

Con PyTorch, la misma idea se escribe de esta manera:

y = torch.relu(x @ W + b)

Un desarrollador no quiere lidiar con registros de GPU, diseño de memoria, programación de instrucciones ni miles de hilos de hardware. Ellos solo indican qué cálculo desean realizar. Python y PyTorch ofrecen una forma sencilla de expresarlo, mientras que otro sistema debe traducir esa descripción a lenguaje de hardware.

El grafo de cálculo

Una red neuronal puede considerarse como un grafo de cálculo. La expresión:

y = relu(x @ W + b)

puede representarse como un gráfico de operaciones conectadas:

Ahora el compilador cuenta con algo concreto en lo que trabajar. Una vez que las operaciones son visibles, puede determinar si se puede acelerar el cálculo. Aquí es donde la optimización adquiere importancia central.

Fusión de operadores y tráfico de memoria

Tomemos un modelo que contenga una cadena como esta:

MatMul
   ↓
Add
   ↓
ReLU

Un compilador puede darse cuenta de que esos pasos pueden combinarse o ajustarse para el hardware objetivo. Una técnica clave es la fusión de operadores: en lugar de tratar cada operación como un paso separado, algunas se combinan, reduciendo el trabajo innecesario, especialmente el traslado de datos entre la memoria y las unidades de cálculo.

Eso es importante porque mover datos es costoso en las cargas de trabajo de inteligencia artificial modernas. El cuello de botella no es solo la velocidad con la que se puede calcular un resultado, sino también la eficiencia con la que los datos de entrada llegan al proceso de cálculo. Esta es una de las razones por las que este campo resulta interesante.

Representaciones intermedias

A continuación viene un concepto familiar de los cursos clásicos de compilación: la representación intermedia (RI). Un compilador tradicional rara vez pasa directamente por esto:

Source Code

a esto:

Machine Code

Pasa por los IR. Los compiladores de IA siguen el mismo principio, pero sus representaciones deben capturar las ideas importantes para el aprendizaje automático. Mientras que un compilador general piensa en:

variables
functions
loops

un compilador de ML también necesita representar elementos como:

tensors
shapes
matrix operations
convolutions
attention
memory layouts

Eso es lo que hace que proyectos como MLIR sean notables. MLIR proporciona la infraestructura necesaria para representar y transformar cálculos en varios niveles de abstracción.

Descenso paso a paso

El compilador avanza gradualmente desde una forma de alto nivel hacia formas más cercanas al hardware de destino. Imagínese bajando por una escalera:

High-level ML operation
          ↓
Tensor operation
          ↓
Lower-level operation
          ↓
Hardware-oriented operation
          ↓
Machine instructions

En cada nivel, las ideas abstractas se vuelven un poco más concretas, hasta que el compilador genera algo que el hardware destino puede utilizar.

El hardware de base

Ahora llegamos a la parte sobre la que nadie piensa al escribir:

model(input)

El cálculo final podría realizarse en una CPU, GPU, NPU, TPU u otro acelerador de IA. A ninguno de ellos le importa que el código original dijera:

y = relu(x @ W + b)

Necesitan instrucciones ejecutables y datos organizados de manera adecuada para ese dispositivo. Una GPU puede requerir código y modelos de ejecución específicos para ella. Una CPU utiliza instrucciones propias y técnicas de vectorización. Un acelerador podría necesitar un backend completamente distinto. Por lo tanto, el mismo modelo de alto nivel puede compilarse de forma diferente según su destino.

¿Y dónde encaja Python?

Python no ha desaparecido y sigue siendo muy útil. Es la interfaz para definir modelos, controlar la ejecución, cargar datos y crear aplicaciones. Simplemente no realiza necesariamente por sí mismo todas las operaciones numéricas. Un modelo mental simplificado:

Esto cambia la forma en que se percibe la inteligencia artificial. Si se mira únicamente desde el lado del modelo, la imagen es breve:

Data → Model → Prediction

Agrega el compilador y hay otra capa intermedia:

Model
  ↓
Compiler
  ↓
Hardware

El modelo indica qué cálculos deben realizarse. El compilador determina cómo representarlos y optimizarlos. El hardware se encarga de ejecutarlos.

Por qué se necesita otra capa

Si ya existen frameworks, ¿por qué añadir una capa de compilador? Porque el hardware de IA es cada vez más especializado. Existen CPUs, GPUs, TPUs, NPUs, aceleradores de inferencia y chips personalizados, cada uno con características diferentes. Un modelo que funciona bien en una arquitectura puede requerir un ajuste muy distinto en otra. El compilador sirve como puente entre ambos lados. Uno de ellos es la cuestión de qué cálculos desea realizar el modelo:

WHAT
The ML model wants to compute

El otro aspecto es cómo un dispositivo específico debe hacerlo de manera eficiente:

HOW
A specific piece of hardware should execute it efficiently

Ese puente sigue ganando importancia.

Una filosofía familiar en un nuevo dominio

Cualquiera que conozca los compiladores tradicionales reconoce la tubería básica:

Source Code
    ↓
Lexer
    ↓
Parser
    ↓
AST
    ↓
IR
    ↓
Optimisation
    ↓
Machine Code

El aprendizaje automático aplica una filosofía similar a los modelos:

ML Model
    ↓
ML Representation
    ↓
ML IR
    ↓
Optimisation
    ↓
Lowering
    ↓
Hardware Code

El dominio es diferente y los problemas son mucho más difíciles, pero la mentalidad se mantiene: representar, transformar, optimizar, reducir, ejecutar.

Una idea natural siguiente es crear un pequeño compilador de IA para aprender cómo funciona. No competiría con LLVM, XLA o MLIR, y el objetivo sería algo lo suficientemente pequeño como para entenderlo, quizás comenzando con un lenguaje de ML simple:

Simple ML Language
        ↓
       Parser
        ↓
      ML AST
        ↓
       ML IR
        ↓
   Simple Optimiser
        ↓
    Lowering
        ↓
   LLVM / Backend

La pila en una sola línea

La IA suele presentarse a través de una interfaz muy sencilla:

model(input)

Bajo esa única llamada se encuentra toda una capa de software y hardware: frameworks, gráficos, representaciones intermedias, compiladores, optimizaciones, backends, hardware y, finalmente, instrucciones de la máquina. Ejecutar un modelo implica poner en marcha todo un ecosistema de compiladores y hardware, siendo Python solo la capa superior.