English · Español

Lab 03 — `torch.compile` sobre grammar MiniGPT + survey distribuido¶

Compilas el grammar MiniGPT con torch.compile, vuelcas el código que genera Inductor, identificas una kernel fusionada y la explicas. Después escribes el survey de 1 página sobre DDP/FSDP/TP/PP — los cuatro patrones distribuidos, sin implementarlos.

Objetivo¶

Ejecuta torch.compile sobre el forward del grammar MiniGPT (o un sustituto mínimo: Linear(64, 600) → softmax), vuelca los kernels generados por Inductor, identifica un kernel fusionado y explica qué hace. Después escribe un README de survey distribuido de 1 página distinguiendo DDP, FSDP, tensor-parallel y pipeline-parallel.

Setup¶

torch >= 2.1. La ruta de compile CPU funciona sin CUDA.
Port a PyTorch de la Fase 17 (el grammar MiniGPT). Si aún no está portado, usa el sustituto mínimo de abajo.

Parte A — Compila el modelo¶

Sustituto mínimo (úsalo si el port a PyTorch de la Fase 17 no está disponible):

import torch
import torch.nn as nn

class TinyHead(nn.Module):
    def __init__(self, d=64, vocab=600):
        super().__init__()
        self.fc = nn.Linear(d, vocab)
    def forward(self, x):
        return torch.softmax(self.fc(x), dim=-1)

torch.manual_seed(42)
model = TinyHead()
model_c = torch.compile(model, mode="default")

x = torch.randn(2, 64)
y1 = model(x)
y2 = model_c(x)
print("compile match:", (y1 - y2).abs().max().item())   # < 1e-6

Parte B — Vuelca la salida de Inductor¶

Pon la variable de entorno antes de importar torch, o establécela via os.environ:

import os
os.environ["TORCH_LOGS"] = "output_code"
os.environ["TORCH_COMPILE_DEBUG"] = "1"   # writes to /tmp/torchinductor_<user>/

import torch
# ... rest of model code

Después de ejecutar:

ls /tmp/torchinductor_$(whoami)/

Deberías ver un árbol de directorios con archivos .py (los wrappers de Python generados) y archivos .cpp o .cu (los kernels generados). En una build sólo CPU, espera C++; con CUDA, espera Triton.

Parte C — Lee un kernel generado¶

Elige el archivo .py más grande en /tmp/torchinductor_<user>/. Tendrá un aspecto parecido a:

# Generated by torch._inductor
triton_poi_fused_softmax_0 = ...  # or cpp_fused_softmax_0 for CPU

@triton.jit  # or extern C
def kernel(...):
    # one or more aten ops fused
    ...

def call(args):
    # orchestration
    ...

Para nuestro modelo Linear + softmax, deberías ver al menos un kernel de softmax fusionado: computa el max, exp, sum, divide en una sola pasada sin materializar el tensor intermedio exp a través de fronteras de kernel.

Guarda el kernel más interesante en experiments/25-compile/kernel.py (o .cpp).

Parte D — Anota el kernel¶

En experiments/25-compile/KERNEL_ANNOTATION.md, recorre el kernel línea por línea. Identifica:

Qué ops de aten se fusionaron. (Probablemente: max, sub, exp, sum, div.)
Desde dónde se lee la entrada. La aritmética de punteros / expresión de índice.
Dónde se escribe la salida.
Si el kernel usa una reducción. Softmax requiere reducciones para max y sum; ¿cómo las expresa Inductor?
Qué no está ahí. Sin búfer intermedio para exp(x - m) — eso es la victoria de la fusión.

Esto es leer, no escribir. No necesitas modificar el kernel. El objetivo es ver que la salida de Inductor es código generado, no magia.

Parte E — Perfila compilado vs eager¶

import time

def bench(fn, x, n=1000):
    # warm-up
    for _ in range(10): fn(x)
    t0 = time.perf_counter()
    for _ in range(n): fn(x)
    return (time.perf_counter() - t0) / n * 1e6   # μs per call

x = torch.randn(2, 64)
print("eager:    ", bench(model, x), "μs")
print("compiled: ", bench(model_c, x), "μs")

En CPU, la versión compilada puede ser 1.1×-2× más rápida, o en nuestro caso minúsculo posiblemente más lenta (el overhead domina para modelos pequeños). Documenta lo que ves — ambos resultados son lecciones válidas.

Parte F — El survey distribuido¶

Escribe experiments/25-compile/DISTRIBUTED.md (~1 página, ~500 palabras). Cuatro secciones, ~125 palabras cada una:

1. DDP — Distributed Data Parallel¶

Patrón: modelo replicado en cada device; trozo de datos repartido entre devices; gradientes all-reduce tras cada .backward().
API: nn.parallel.DistributedDataParallel(model).
Cuándo usar: el modelo cabe en un device, y estás escalando rendimiento.
Comunicación: un all-reduce por parámetro, por paso. Solapado con el backward.
Cuándo NO usar: el modelo no cabe. Usa FSDP o TP.

2. FSDP — Fully Sharded Data Parallel¶

Patrón: parámetros, gradientes y estado del optimizador están shardeados entre devices. Cada device guarda 1/N. Durante el forward, los parámetros de la capa se all_gatheran desde los pares; se liberan tras la capa. Backward similar.
API: torch.distributed.fsdp.FullyShardedDataParallel(model).
Cuándo usar: el modelo no cabe en un device, pero una capa sí.
Comunicación: all-gather por capa del forward + reduce-scatter por capa del backward. Mucha más comunicación que DDP.
Cuándo NO usar: el modelo cabe en un device (DDP es más barato) o incluso una capa no cabe (necesitas TP).

3. Tensor parallel (TP)¶

Patrón: dentro de una sola capa, la matriz de pesos se reparte entre devices (por filas o columnas). El matmul se particiona; las salidas se concatenan.
API: a nivel de librería (Megatron-LM, FairScale, torch.distributed.tensor.parallel).
Cuándo usar: los pesos de una sola capa no caben. Común en modelos 70B+ para el LM head.
Comunicación: un all-reduce por capa (para la forma row-split). Alta; requiere interconexión de calidad NVLink.
Cuándo NO usar: la comunicación es lenta relativa al cómputo. Usa FSDP en su lugar.

4. Pipeline parallel (PP)¶

Patrón: modelo repartido a lo largo de la profundidad. Device 0 guarda las capas 1-10, device 1 las capas 11-20, etc. Las activaciones fluyen hacia adelante, los gradientes hacia atrás, en un patrón de burbuja.
API: torch.distributed.pipeline.sync.Pipe o wrappers de librería.
Cuándo usar: el modelo tiene muchas capas secuenciales; el ancho de banda de comunicación entre devices es bajo.
Comunicación: un send/recv por micro-batch por etapa. Volumen bajo pero latencia alta.
Cuándo NO usar: pocas capas, mucho ancho de banda — DDP o TP dominan.

Termina el survey con un párrafo de 3 frases "cuál elegiría" para el grammar MiniGPT (respuesta: DDP, porque el modelo es minúsculo — pero la pregunta busca hacerte razonar).

Parte G — Escribe el informe¶

experiments/25-compile/REPORT.md:

La comprobación de coincidencia de compile (Parte A): error máximo < 1e-6.
Puntero al archivo de kernel (Parte C) y la anotación (Parte D).
Números de perfilado (Parte E) con interpretación honesta (los modelos pequeños pueden no beneficiarse).
Puntero a DISTRIBUTED.md (Parte F).

Entregable¶

experiments/25-compile/: - REPORT.md. - kernel.py o kernel.cpp — el kernel generado por Inductor. - KERNEL_ANNOTATION.md — tu walkthrough. - DISTRIBUTED.md — el survey de 1 página. - manifest.json.

Aceptación¶

El modelo torch.compile'd coincide con la salida eager dentro de 1e-6.
Un kernel generado por Inductor está guardado y anotado.
La anotación identifica correctamente el patrón de softmax fusionado.
El survey distribuido distingue los cuatro patrones en forma de 2 frases por patrón.

Pitfalls¶

TORCH_LOGS puesto después del import. La variable de entorno debe ponerse antes de import torch. Lo más fácil: ponla en tu shell o usa os.environ al principio del script antes de cualquier import de torch.
/tmp/torchinductor_<user>/ limpiado entre ejecuciones. Inductor cachea por hash del grafo; limpiar el directorio fuerza un compile fresco. Útil para debugging.
La primera llamada es lenta. torch.compile es JIT — la primera llamada traza y compila. Bencharca sólo la ruta caliente.
Recompilaciones por cambio de forma. Pasa la misma forma a cada llamada al benchmarkear, o usa mode="reduce-overhead" con cuidado.
El compile en CPU usa C++. Espera archivos .cpp y .so, no Triton. La lección es la misma — fusionado, generado, legible.
"Mi MiniGPT aún no está portado a PyTorch." Usa el sustituto. El punto es leer la salida de Inductor, lo cual no depende del tamaño del modelo.

Stretch¶

Compila el grammar MiniGPT completo (bloque decoder + LM head). Identifica un kernel de attention+softmax fusionado. Compara con la vista previa de flash-attention de la Fase 27.
Ejecuta con mode="max-autotune". Compara tiempo de compile y runtime con mode="default".
Usa torch._dynamo.export para extraer el FX graph como artefacto autónomo.

Fin de los laboratorios de la Fase 25. Toca escribir PHASE_25_REPORT.md y prepararse para la Fase 26.

Siguiente: Fase 26 — Cuantización.

Lab 03 — torch.compile sobre grammar MiniGPT + survey distribuido¶