Entrenar más rápido

4.1Warp y SIMT

Las GPU de NVIDIA ejecutan los hilos en grupos de 32 llamados warps. Es Single Instruction, Multiple Threads (SIMT): el warp comparte un flujo de instrucciones y cada hilo tiene sus propios registros.

Dos consecuencias dominan el rendimiento de los kernels:

1. Divergencia de ramas: si los hilos de un warp toman caminos distintos en un if, el warp ejecuta ambos caminos con carriles enmascarados. Mantén el control de flujo uniforme dentro del warp cuando puedas.
2. Coalescencia de memoria: cuando los 32 hilos leen direcciones contiguas, el hardware las une en unas pocas transacciones anchas. Los accesos dispersos multiplican el tráfico de memoria. Ver coalescencia de memoria.

Los warps también son la unidad que el planificador hace malabares para ocultar la latencia. Más warps residentes (mayor ocupación) le dan más opciones, hasta que se agotan los registros y la memoria compartida.

4.2Coalescencia de memoria

La memoria global se accede en segmentos alineados (sectores de 32 bytes, agrupados en líneas de 128 bytes). Si los 32 hilos de un Warp y SIMT leen 32 floats consecutivos de 4 bytes, el hardware los atiende con una transacción de 128 bytes. Si leen con un salto grande, puede necesitar 32 transacciones separadas, una penalización de 32× en ancho de banda.

Arreglos clásicos de la práctica de programación de GPU (p. ej. la guía de PRACE)

• Estructura de arreglos en lugar de arreglo de estructuras.
• Que el hilo i toque el elemento i, no i × salto.
• Pasar los patrones de acceso incómodos (transposiciones, stencils) por memoria compartida y luego escribir de forma coalescente.
• Cuidado con los conflictos de bancos de memoria compartida al hacerlo.

Los profilers (Nsight Compute) reportan "sectores por solicitud". Lo ideal es 4 para cargas de 4 bytes.

4.3Ocupación

Cada SM puede alojar un número limitado de warps (64 en la H100). La ocupación es cuántos están realmente residentes, limitada por el uso de cada bloque de:

• Registros por hilo (el banco de registros lo comparten todos los warps residentes),
• Memoria compartida por bloque,
• Hilos por bloque y bloques por SM.

Mayor ocupación le da al planificador de warps más candidatos a los que cambiar mientras otros esperan a la memoria. Pero es un medio, no un fin. Muchos kernels rápidos (GEMM grandes, FlashAttention) corren con ocupación baja y usan registros y memoria compartida para reutilizar datos.

Enfoque práctico (según las buenas prácticas de CUDA y PRACE): usa la calculadora de ocupación o Nsight Compute para ver el recurso limitante, prueba algunos tamaños de bloque (múltiplos de 32) y mide. No supongas.

4.4Tensor Core

Los tensor cores calculan productos de matrices pequeñas (p. ej. bloques de 16×16) en una sola instrucción. En una H100 entregan ~989 TFLOP/s densos en BF16 frente a ~67 TFLOP/s en FP32 simple. Esa diferencia es la razón por la que el entrenamiento en precisión mixta es lo normal.

Cómo obtener los FLOPs

• Usa BF16/FP16 (o FP8 con Transformer Engine) para las multiplicaciones de matrices. FP32 sin TF32 nunca los usa.
• Mantén las dimensiones de las GEMM grandes y múltiplos de 8/16. Micro-batches diminutos o shards de paralelismo tensorial muy estrechos generan GEMM pequeñas que no saturan las unidades.
• Mantén los datos fluyendo desde HBM. Los tensor cores son tan rápidos que muchas operaciones (normalización, activaciones, softmax de atención) pasan a estar limitadas por memoria. Para eso existe la fusión de kernels como FlashAttention.

El Model FLOPs Utilization (MFU (Model FLOPs Utilization)) es, en el fondo, una medida de qué tan ocupados están los tensor cores.

4.5Precisión mixta (BF16, FP8)

El entrenamiento en precisión mixta ejecuta el forward y el backward en un formato de baja precisión y mantiene pesos maestros en FP32 para la actualización del optimizador:

• BF16 conserva el rango de exponente de FP32, así que rara vez necesita loss scaling. Es lo predeterminado en A100/H100.
• FP16 necesita loss scaling dinámico para que los gradientes pequeños no se vuelvan cero.
• FP8 (H100 en adelante) vuelve a duplicar aproximadamente el rendimiento de las matmul. Bibliotecas como Transformer Engine gestionan el escalado por tensor.

Por qué es tan importante: los tensor cores entregan ~15× más FLOP/s en BF16 que en FP32 simple en una H100. Entrenar en FP32 deja casi toda la GPU ociosa.

Memoria: pesos y gradientes en 2 bytes, más 12 bytes de estado del optimizador en FP32 (maestro, m y v de Adam), da la regla de "16 bytes por parámetro" que se usa en todo LazyLab.

4.6Paralelismo de datos (DDP)

La forma más simple de escalar: replica el modelo en N GPUs, dale a cada una una parte distinta del batch y luego haz all-reduce de los gradientes para que cada réplica dé el mismo paso del optimizador.

Costo: un all-reduce del gradiente completo (2 bytes × parámetros en BF16) por paso. Con all-reduce en anillo, el tráfico por GPU es ~2× el tamaño del gradiente, independiente de N. Por eso el paralelismo de datos escala tan bien.

Cómo hacerlo rápido

• Solapamiento: DDP agrupa los gradientes en buckets y hace all-reduce de cada bucket en cuanto el backward lo produce, ocultando la comunicación detrás del cómputo (ver CUDA streams y ejecución asíncrona).
• Acumulación de gradientes: corre varios micro-batches antes de sincronizar. Obtienes un batch global mayor con la misma memoria y menos sincronizaciones por token.

Límite: cada GPU debe tener el estado completo del modelo. Cuando eso deja de caber, ver ZeRO y FSDP, paralelismo tensorial y paralelismo de pipeline.

4.7CUDA streams y ejecución asíncrona

Un stream es una cola ordenada de operaciones de GPU. El trabajo en streams distintos puede correr en paralelo: una copia host→dispositivo en uno, un kernel en otro, una colectiva de NCCL en un tercero. Los eventos expresan dependencias entre streams sin bloquear la CPU.

Este es el mecanismo detrás de cada optimización de "solapamiento":

• El all-reduce de gradientes que se solapa con el backward en DDP usa un stream de comunicación aparte.
• Precargar el siguiente batch mientras se calcula el actual (ver pipeline de entrada).
• ZeRO-3/FSDP precargando los parámetros de la siguiente capa (ver ZeRO y FSDP).

Proyectos como grCUDA van más lejos. Construyen el grafo de dependencias a partir de los datos que toca cada kernel y programan kernels independientes en streams separados automáticamente, así los programas políglotas obtienen concurrencia sin escribir código de streams a mano.

Trampa: cualquier sincronización implícita (cudaDeviceSynchronize, .item() en PyTorch, copias desde memoria paginable) colapsa tu solapamiento cuidadosamente construido en ejecución serial.

4.8Checkpointing de activaciones

La retropropagación necesita las activaciones del forward. Para una capa transformer con FlashAttention son aproximadamente 34 · s · b · h bytes (secuencia × micro-batch × hidden), y se acumulan rápido. Un modelo de 7B con micro-batch 4 necesita más de 70 GB de activaciones.

El checkpointing (también llamado recomputación) guarda solo la entrada de cada capa y vuelve a ejecutar el forward de la capa durante el backward. Costo: un forward extra, ~33% más de cómputo (8N en lugar de 6N FLOPs por token). Beneficio: la memoria de activaciones baja a unos 2 · s · b · h bytes por capa.

Úsalo cuando las activaciones son lo que te impide hacer caber el modelo o un micro-batch mayor. Apágalo cuando la memoria alcanza. Entonces es puro costo. El checkpointing selectivo (solo el interior de la atención) logra la mayor parte del ahorro por una fracción del costo.

Nota que el MFU deliberadamente no cuenta los FLOPs recalculados, así que el checkpointing baja el MFU aunque la GPU esté igual de ocupada.

4.9ZeRO y FSDP

El paralelismo de datos simple desperdicia memoria: cada rango guarda copias idénticas de todo. ZeRO (y FSDP de PyTorch) lo reparten:

El problema de la etapa 3: los parámetros deben reunirse con all-gather antes del forward y del backward de cada capa, en cada micro-batch. Para un modelo de 70B son ~140 GB movidos dos veces por micro-batch. Sobre NVLink va bien. Sobre una NIC de 400 Gb/s puede dominar el paso.

Guía práctica

• ZeRO-1 es casi gratis y casi siempre vale la pena.
• ZeRO-3/FSDP es excelente dentro de un nodo o en redes rápidas, y para hacer caber un modelo que de otro modo no cabría.
• A gran escala, combínalo: TP dentro del nodo, ZeRO-1 o sharding híbrido (repartir dentro del nodo, replicar entre nodos).

4.10Paralelismo tensorial

El paralelismo tensorial (TP) reparte las matrices de pesos dentro de cada capa. En el esquema de Megatron-LM, la primera matmul del MLP se divide por columnas y la segunda por filas, así que una capa solo necesita un all-reduce en el forward y uno en el backward. Con la atención tratada igual, son ~4 all-reduce de activaciones por capa transformer por micro-batch.

Compromisos

• ✅ Pesos, gradientes, estado del optimizador y activaciones se reducen por el grado de TP.
• ❌ La comunicación es frecuente, sensible a la latencia y está en el camino crítico. Casi no se puede ocultar.
• ❌ Los shards se vuelven estrechos: hidden/TP demasiado pequeño significa GEMM pequeñas y menor eficiencia de los tensor cores.

Regla: mantén el TP dentro del dominio NVLink (TP ≤ 8 en servidores HGX). TP=16 a través de InfiniBand y RoCE es una forma clásica de desperdiciar un clúster.

4.11Paralelismo de pipeline

El paralelismo de pipeline (PP) asigna capas contiguas a etapas. La etapa 1 corre las capas 1–20, la etapa 2 las 21–40, y así. Solo las activaciones en los bordes entre etapas cruzan la red, así que PP se lleva bien con varios nodos sobre InfiniBand y RoCE.

La burbuja: al inicio y al final de cada paso algunas etapas quedan ociosas mientras el pipeline se llena y se vacía. Con p etapas y m micro-batches por paso (calendario 1F1B):

fracción de burbuja ≈ (p − 1) / (m + p − 1)

Más micro-batches (acumulación de gradientes) achican la burbuja. Los calendarios intercalados la reducen más, a cambio de más comunicación.

Restricciones: el número de capas debe dividirse (casi) por igual entre etapas, y las etapas deben estar equilibradas. La primera etapa además guarda la mayor cantidad de activaciones en vuelo.

La receta clásica para modelos grandes, el paralelismo 3D, combina los tres: TP dentro del nodo, PP entre unos pocos nodos y DP (a menudo con ZeRO-1) en todo lo demás.