English · Español

Lab 02 — Inferencia tensor parallel en 2 GPUs cloud (el único lab con gasto)¶

Objetivo: ejecutar MiniGPT-grammar con tensor parallel entre 2 GPUs cloud. Ver el all-reduce en NCCL. Medir el speedup (¡espera slowdown!). Gasta ≤ $3.

Tiempo estimado: 3–4 horas de wall clock, ≤ 3 horas de runtime real de GPU.

Prerrequisito: Labs 00 y 01 completos. BudgetGuard funciona. Cuenta del proveedor financiada con el tope del lab (p. ej., $3 prepagados en RunPod). Checkpoint entrenado de MiniGPT-grammar (Fase 17) listo para subir.

El único lab en todo el currículo que gasta dinero real. Lee la matemática de ancho de banda de la teoría 03 y la checklist del lab 00 antes de levantar nada.

Qué produces¶

Una corrida funcional de inferencia TP en experiments/35-tp-inference/ que:

Usa el BudgetGuard del lab 00 para autorizar el spinup.
Levanta una instancia de 2× GPU tier consumer (RTX 4090 o similar, single-node, idealmente con NVLink).
Sube el checkpoint entrenado de MiniGPT-grammar y el conjunto de prompts de test ("he goed to school" etc.).
Ejecuta un bucle de inferencia TP en 2 GPUs usando el backend NCCL, haciendo shard de la tabla de embeddings entre los 2 workers.
Compara latencia por token vs single-GPU. Registra el ratio.
Termina la instancia. Screenshot de la consola del proveedor prueba la terminación.
Loguea el $ real gastado en experiments/35-cloud-budget/spend.jsonl vía BudgetGuard.record_actual.

Más una extensión continuada a src/minitrain/:

src/minitrain/tensor_parallel.py — clases ColumnParallelLinear, RowParallelLinear, ShardedEmbedding. Implementaciones educativas siguiendo el patrón de Megatron (theory/02). No de calidad de producción.

TODOs¶

Bloque A — implementa las primitivas TP (primero localmente)¶

Antes de cualquier spinup cloud, escribe y testea src/minitrain/tensor_parallel.py en un entorno mock de single-CPU world_size=2 usando torch.multiprocessing.spawn:

# src/minitrain/tensor_parallel.py — esqueleto (Borja escribe el cuerpo)

import torch
import torch.distributed as dist
from torch.nn import Module, Parameter

class ColumnParallelLinear(Module):
    """Linear(in_features, out_features // world_size) per worker.
    Output is sharded along out_features; *no* comm at the output.
    Input is replicated (must match) — *no* comm at the input.
    """
    def __init__(self, in_features: int, out_features: int): ...
    def forward(self, x): ...

class RowParallelLinear(Module):
    """Linear(in_features // world_size, out_features) per worker.
    Input is sharded along in_features; output is full but partial-sum.
    Forward includes one all-reduce.
    """
    def __init__(self, in_features: int, out_features: int): ...
    def forward(self, x): ...

class ShardedEmbedding(Module):
    """nn.Embedding sharded by vocab_size // world_size.
    Lookup is local-only for hits in the worker's slice; an all-reduce
    combines partial outputs.
    """
    def __init__(self, vocab_size: int, embedding_dim: int): ...
    def forward(self, ids): ...

Tests en tests/minitrain/test_tensor_parallel.py:

test_column_parallel_matches_single_gpu — la salida de la linear TP (tras concatenar shards) coincide con la linear no-TP dentro de 1e-4.
test_row_parallel_matches_single_gpu — igual.
test_sharded_embedding_lookup — la salida del embedding con shard coincide con la no-sharded.
test_mlp_pattern_minimal — ColumnParallel → GELU → RowParallel con un all-reduce coincide con el MLP single-GPU.

Estos tests corren en la CPU local antes de cualquier gasto cloud.

Bloque B — prepara el script del experimento cloud¶

experiments/35-tp-inference/run.py:

# Skeleton — Borja writes the body

def main():
    rank, world_size = init_distributed("nccl")
    assert world_size == 2, "This lab is 2-GPU only."

    model = build_minigpt_grammar_tp(rank, world_size)   # TP-sharded version
    load_checkpoint(model, "checkpoint-trained.pt", strict_shard_aware=True)

    prompts = ["he goed to school", "i has eat", "she dont like apples", ...]

    times_single = run_single_gpu_baseline(prompts) if rank == 0 else None
    dist.barrier()

    times_tp = run_tp_inference(model, prompts)

    if rank == 0:
        save_results(times_single, times_tp, "results.json")

    cleanup()

El baseline single-GPU corre solo en el rank 0 (el otro rank queda ocioso) para que la comparación sea manzanas con manzanas.

Bloque C — pre-flight (obligatorio)¶

Antes de runpodctl create:

BudgetGuard.remaining lee ≥ $3.50 (tope + buffer).
Consola del proveedor: alerta de presupuesto al 80% puesta; auto-terminación a 4 horas puesta.
BudgetGuard.authorize("lab02-tp-spinup-rtx4090x2", 3.00) retorna sin levantar.
Captura el ID de instancia + timestamp de inicio en experiments/35-tp-inference/instance.json.

Bloque D — ejecuta¶

runpodctl create ... (o equivalente del proveedor). Espera a que la instancia esté lista.
scp del checkpoint + script arriba.
ssh <instance> "torchrun --nproc-per-node=2 experiments/35-tp-inference/run.py".
Recupera results.json de vuelta.
runpodctl terminate <instance-id>. Verifica terminación en la consola del proveedor; screenshot.
BudgetGuard.record_actual("lab02-tp-spinup-rtx4090x2", <actual $>) usando el real reportado por la consola del proveedor.

Bloque E — analiza¶

Calcula y registra en manifest.json:

Latencia por token single-GPU (media sobre 100 tokens × 10 prompts).
Latencia por token 2-GPU TP.
Speedup = single / TP (esperado: <1, es decir, slowdown).
Wall time del all-reduce por token (desde dentro del bucle TP).
Ratio comm/cómputo.

Escribe un párrafo de reflexión de "por qué 2-GPU TP es más lento que single-GPU aquí" en el manifest. Razonamiento esperado: el modelo es diminuto (~500k params), el volumen del all-reduce por token es pequeño pero no nulo, NVLink ayuda pero no puede hacer que una corrida TP de 2 GPUs sea más rápida que una corrida de 1 GPU para un modelo que no estaba memory-bound en una sola GPU. La lección generaliza: TP te cuesta latencia a pequeña escala; te compra margen a gran escala.

Bloque F — curva de scaling (sintética)¶

El slowdown de 2-GPU TP es esperado para el tutor de gramática. Para hacer la lección concreta, genera el gráfico prospectivo:

Calcula la latencia single-GPU predicha para modelos hipotéticos con $d_{\text{model}} = 64, 256, 1024, 4096$ con el mismo patrón TP.
Calcula la latencia 2-GPU TP predicha usando la matemática de ancho de banda del all-reduce de theory/03.
Grafica el cruce: ¿a qué $d_{\text{model}}$ TP se vuelve más rápido?

Commitea experiments/35-tp-inference/scaling-curve.png con el cruce anotado. Esperado: el cruce está alrededor de $d_{\text{model}} = 1024$–$2048$ en el hardware de test.

Bloque G — manifest¶

experiments/35-tp-inference/manifest.json:

{
  "seed": 35200,
  "lab": "02-tp-inference-cloud",
  "vendor": "<chosen>",
  "instance_type": "<e.g., 2x RTX 4090>",
  "nvlink_present": true,
  "instance_id": "<filled in>",
  "start_ts": "<filled in>",
  "end_ts": "<filled in>",
  "duration_hours": "<filled in>",
  "estimated_cost_usd": 3.00,
  "actual_cost_usd": "<filled in>",
  "terminated_screenshot": "experiments/35-tp-inference/proof-terminated.png",
  "single_gpu_per_token_ms": "<filled in>",
  "tp_per_token_ms": "<filled in>",
  "speedup": "<filled in: should be < 1>",
  "lesson_notes": "<the why-slower paragraph>",
  "crossover_d_model": "<filled in from scaling-curve>"
}

Restricciones¶

Tope duro de $3. BudgetGuard.authorize lo fuerza. Más allá de eso, BudgetGuardExceeded se levanta — no la captures.
Solo tier spot. Si on-demand es la única opción, la matemática del presupuesto no funciona — no ejecutes el lab; documenta la restricción. Algunas semanas el precio se moverá y el lab genuinamente no es ejecutable con $3. Es un resultado aceptable — escribe el análisis sin la corrida cloud, marcado como EDUCATIONAL_STUB.
Single-node, 2 GPUs. Multi-nodo está fuera de alcance.
Sin tuning de NCCL. Nada de NCCL_DEBUG, nada de NCCL_IB_DISABLE. Por defecto. El objetivo es ver la comm, no optimizarla.
Termina antes de commitear. Si el lab falla a mitad, termina la instancia primero, luego depura localmente.
Tope de wall clock de dos horas. Más allá de 2 horas de runtime de GPU, termina y analiza lo que tengas. La tasa coste-vs-aprendizaje ya ha tocado techo.

Condiciones de parada¶

Has acabado cuando:

experiments/35-tp-inference/manifest.json tiene actual_cost_usd rellenado y ≤ $3.00.
proof-terminated.png muestra la consola del proveedor con la instancia terminada.
BudgetGuard.total_spent coincide con el coste real del manifest (dentro del redondeo).
El párrafo de "por qué más lento" y el gráfico de la curva de cruce están commiteados.
Los tests de src/minitrain/tensor_parallel.py pasan en el mock de la CPU local.

Pista de último recurso¶

Si algo va mal: termina primero, depura después. Una instancia 2-GPU olvidada quema ~$0.70/hr = ~$17/día. Si te despiertas con "¿terminé?", termina ahora, comprueba luego.

Si la corrida cloud se cuelga y no puedes hacer ssh para terminar: usa la consola web del proveedor para terminar. Marca la página de terminate como favorito antes de empezar.

Si el análisis del lab sale ininteligible (p. ej., la corrida 2-GPU misteriosamente más rápida que la single-GPU en un modelo de 500k params — eso sería no-físico): tu baseline single-GPU probablemente no era realmente single-GPU. Comprueba que CUDA_VISIBLE_DEVICES=0 estaba puesto para el baseline.

Cuándo consultar `solutions/`¶

Tras commitear el experimento. La solución vive en solutions/02-tp-inference-ref.md — escrita en la apertura de la fase tras tener los internos de PyTorch de la Fase 25 de Borja. La solución intencionalmente no incluye cifras de coste cloud reales — esas dependen del precio del proveedor el día de la ejecución; la solución explica la forma esperada (TP más lento que single-GPU a esta escala) y la matemática del cruce, no la cantidad en dólares.

Siguiente lab: lab/03-megatron-fsdp-reading.md.