English · Español

03 — `torch.compile` y Distributed (Survey)¶

Esta página cubre dos temas que merecerían cada uno una fase completa pero que en Fase 25 son un survey honesto: cómo funciona torch.compile (Dynamo → AOTAutograd → Inductor) y cuáles son los cuatro patrones canónicos de distributed (DDP, FSDP, tensor-parallel, pipeline-parallel). Hands-on con compile, lectura para distributed. Phase 33 vuelve a compile; Phase 35 hace distributed real.

Esta página es intencionalmente un survey. El laboratorio de la Fase 25 usará torch.compile (lo ejecutará, volcará la salida de Inductor, identificará la fusión); no construirá un pipeline de compile. Distributed es de sólo lectura: el laboratorio escribe un README de 1 página distinguiendo los cuatro patrones, sin código de torch.distributed más allá del hola-mundo de init_process_group.

El hands-on completo de compile es la Fase 33 (serving) y el de distributed es la Fase 35.

Parte A: `torch.compile`¶

Qué es¶

torch.compile(model) devuelve una versión optimizada del modelo. Las llamadas posteriores trazan el forward (y el backward) en un grafo, optimizan el grafo (fusión, elección de layout, selección de kernel), emiten kernels Triton (para GPU) o C++ (para CPU), y ejecutan esos kernels en lugar de la secuencia de dispatch por op.

Esquemáticamente:

model(x)                                 # eager: 100 dispatches per forward
↓
model_c = torch.compile(model)
model_c(x)
↓
[TorchDynamo]      Python bytecode → FX graph
[AOTAutograd]      FX graph → joint forward+backward FX graph
[Inductor]         FX graph → Triton/C++ kernel files
[runtime]          Loads the compiled kernels; runs them in place of eager

Tras la primera llamada (tiempo de compile: ~segundos-a-minutos), las llamadas posteriores ejecutan los kernels compilados — normalmente 1.5–3× más rápido que eager para inferencia, 1.2–2× más rápido para entrenamiento.

Etapa 1: TorchDynamo (Python → FX)¶

Dynamo es un tracer de bytecode de Python. Ejecuta el forward() de tu modelo una vez simbólicamente — propagando tensores fake con metadata de forma/dtype — y registra cada operación de torch en un FX graph (la representación intermedia de PyTorch).

Si Dynamo no puede trazar una parte del código (por ejemplo, una rama de Python dependiente del valor de un tensor), inserta un graph break: emite un grafo para el prefijo, ejecuta el Python ofensivo en modo eager, después traza el sufijo. Los graph breaks reducen las oportunidades de optimización.

Causas comunes de graph break:

if x.sum() > 0: (control de flujo de Python sobre tensor).
Llamadas a librerías no trazables.
print(...) con un argumento tensor.
Mutar estructuras de datos de Python.

Diagnóstico verboso de graph breaks: TORCHDYNAMO_VERBOSE=1 python script.py.

Etapa 2: AOTAutograd (FX conjunto forward+backward)¶

El FX graph del forward generado por Dynamo se alimenta a AOTAutograd, que:

Traza el pase backward simbólicamente (igual que haría el motor de autograd en runtime, pero por adelantado).
Produce un FX graph conjunto con nodos tanto de forward como de backward.
Descompone ops de alto nivel (como linear) en sus ops primitivas (matmul, add) para una optimización más fina por Inductor.

Para entrenamiento, aquí es donde el grafo de autograd se "fusiona con" el grafo del forward — ahorrando el coste de captura del grafo por paso.

Etapa 3: Inductor (FX → Triton/C++)¶

Inductor es el backend compilador de PyTorch. Toma el FX graph y emite código de kernel real:

Ruta CUDA: emite kernels Triton (el mismo lenguaje Triton de la Fase 24).
Ruta CPU: emite C++ con OpenMP / intrínsecos vectorizados.

Inductor hace:

Fusión: combina ops elementwise + reducciones adyacentes en un solo kernel.
Selección de layout: elige el layout de memoria (contiguo vs strided) por op.
Loop tiling: elige tamaños de tile para SMEM / cache.
Autotuning: opcionalmente barre tamaños de tile (mode="max-autotune").

La salida se guarda en /tmp/torchinductor_<user>/<hash>/. Con TORCH_LOGS=output_code, Inductor imprime el/los kernel(s) generado(s) a stderr. Leer esto es revelador — es simplemente un archivo Triton, como el que Borja escribió a mano en la Fase 24.

Modos de compile¶

torch.compile(model, mode="default")           # balanced; ~1 minute compile
torch.compile(model, mode="reduce-overhead")   # uses CUDA Graphs; lowest latency
torch.compile(model, mode="max-autotune")      # exhaustive tile-size sweep; slow compile, fastest run

Para inferencia del grammar MiniGPT: reduce-overhead suele ser lo adecuado. El laboratorio 03 prueba cada uno.

Qué genera Inductor para `nn.Linear + softmax`¶

A mano escribirías:

y = lm_head(x)         # nn.Linear(64, 600) → cuBLAS gemm
p = F.softmax(y, -1)   # custom softmax kernel

Tras torch.compile, Inductor podría emitir:

Una llamada cuBLAS gemm para el matmul (no suele fusionar matmuls con elementwise).
Un kernel Triton fusionado para el softmax (max + exp + sum + normalize, todo en un solo lanzamiento).

Para el laboratorio, vuelca esto con TORCH_LOGS=output_code e identifica la fusión de softmax en la fuente de Inductor. El kernel son ~30 líneas de Triton — directamente comparable al que Borja escribió a mano en la Fase 24.

Cuándo ayuda el compile y cuándo no¶

Caso	Speedup
Muchas ops elementwise pequeñas entre matmuls	2–5× (la fusión elimina tensores intermedios)
Un matmul grande domina el tiempo	~1× (cuBLAS ya es óptimo)
Graph breaks cada pocas ops	Marginal (sólo overhead)
Modelo con `torch.jit.script` ya aplicado	Posiblemente negativo (compile re-traza, puede regresar)

El laboratorio de serving de la Fase 33 medirá las ganancias de compile sobre el grammar MiniGPT completo.

Parte B: Distributed (Survey)¶

Esto es un survey de conceptos. La Fase 35 construye estos de verdad. Aquí los nombramos, los describimos y los situamos en un eje 2D de qué se reparte y cómo escala la comunicación.

DDP: Data-Parallel¶

GPU 0:  model (full copy) + batch slice 0
GPU 1:  model (full copy) + batch slice 1
GPU 2:  model (full copy) + batch slice 2
GPU 3:  model (full copy) + batch slice 3

Cada GPU mantiene una copia completa del modelo. Diferentes GPUs ven diferentes trozos del batch. El forward y el backward son independientes por GPU; tras el backward, los gradientes se hacen all-reduce entre GPUs (se promedian) para que todas las copias se mantengan sincronizadas.

PyTorch: torch.nn.parallel.DistributedDataParallel(model).

Comunicación: O(tamaño del modelo) por paso (el all-reduce de gradientes). Escala bien hasta 8 GPUs; sufre más allá de 64 porque cada GPU sigue guardando el modelo completo.

FSDP: Fully-Sharded Data-Parallel¶

Igual que DDP, pero cada GPU mantiene sólo un shard de los parámetros de cada capa. Antes de ejecutar una capa, el shard se junta desde las otras GPUs (allgather); tras la capa, el shard se libera.

Reduce la memoria por GPU en N× (donde N es el world size).
Aumenta la comunicación (allgather + reduce-scatter por capa frente a un all-reduce por paso).
Necesario para modelos que no caben en una sola GPU.

PyTorch: torch.distributed.fsdp.FullyShardedDataParallel.

Tensor-Parallel (TP)¶

Un único matmul se reparte entre GPUs. Para \(Y = X W\) con \(W\) una matriz 4096×4096 en 2 GPUs:

GPU 0:  W_left  (4096 × 2048)
GPU 1:  W_right (4096 × 2048)

X is broadcast to both GPUs.
Y_left  = X @ W_left   on GPU 0   (output shape (B, 2048))
Y_right = X @ W_right  on GPU 1   (output shape (B, 2048))
Y = concat(Y_left, Y_right)

Para attention: reparte heads entre GPUs. Para FFN: reparte la dim oculta.

Comunicación: por capa (cada forward a través de una capa partida con TP requiere un allreduce o allgather de las salidas parciales).

Librería: Megatron-LM, vLLM, o escrita a mano. El tensor_parallel de PyTorch está a nivel alpha.

Pipeline-Parallel (PP)¶

Reparte la profundidad del modelo. Capas 1–8 en GPU 0; capas 9–16 en GPU 1; capas 17–24 en GPU 2; capas 25–32 en GPU 3.

Para un forward: GPU 0 ejecuta las capas 1–8, pasa activaciones a GPU 1, ejecuta las capas 9–16, etc. PP ingenuo infrautiliza las GPUs (sólo una está activa a la vez). El arreglo: micro-batching — repartir un batch en K micro-batches, hacer que GPU 0 procese el micro-batch 2 mientras GPU 1 procesa el micro-batch 1, etc. Éste es el "schedule de pipeline" / "burbuja" del paralelismo por pipeline.

Librerías: pippy de PyTorch, pipeline de DeepSpeed.

Elegir entre los cuatro¶

¿Cabe el modelo en 1 GPU?	Patrón
Sí	DDP (más simple, escala hasta unas 8 GPUs)
No, pero por capa cabe	FSDP (sharding de params) o TP (reparte los matmuls)
No, y por capa tampoco cabe	PP (reparte la profundidad) o 3D-parallel (combinar FSDP + TP + PP)

Para el laboratorio de la Fase 35, el grammar MiniGPT cabe en el portátil de Borja — distributed no está motivado por necesidad. La Fase 35 usa un modelo algo mayor (o simula multi-GPU con backend gloo en CPU) para demostrar los patrones.

Lo que deberías ahora ser capaz de hacer¶

Ejecutar torch.compile(model) y volcar la salida de Inductor.
Leer un kernel Triton generado por Inductor e identificar la fusión.
Distinguir DDP, FSDP, TP, PP — qué se reparte, qué se comunica.
Predecir qué patrón distribuido aplica a un tamaño de modelo y conteo de GPUs dados.
Reconocer las limitaciones: graph breaks de compile, coste de comunicación de FSDP, overhead de latencia de TP, burbuja de PP.

Lo que esta página NO cubre¶

Fallos de compile y sus arreglos. La Fase 33 dedica tiempo a debugéar compile.
CUDA Graphs. Fase 33.
Algoritmos de scheduling de pipeline (1F1B, 1F1B intercalado). Fase 35.
Primitivas colectivas NCCL en detalle. Fase 35.
Transformaciones funcionales torch.func. La Fase 38 quizá.

Siguiente: lab/00-dispatcher-trace.md — instrumenta una llamada linear(x, W, b) y lee el log del dispatcher.

03 — torch.compile y Distributed (Survey)¶

Parte A: torch.compile¶