Escalado de LLMs con paralelismo multi-GPU y multinodo

Viacheslav Dubrov

2025-05-04 · Actualizado 2026-07-15 · 11 min · Ingeniería de modelos · Infraestructura de IA

Traducción automática Este artículo se tradujo automáticamente a partir de la versión original en inglés.

Las cargas de trabajo con modelos grandes superan los límites de una sola GPU por distintos motivos. Un trabajo de entrenamiento puede quedarse sin memoria por el estado del optimizador. Otro puede hacerlo por las activaciones de secuencias largas. Aunque el modelo quepa, quizá no alcance el throughput objetivo. Cada problema requiere una partición y un patrón de comunicación diferentes.

Este es un recorrido práctico por las principales estrategias de paralelismo y las restricciones que las condicionan, basado en Ultra-Scale Playbook de Hugging Face. El objetivo es mostrar qué aporta cada partición, qué comunica y cuándo resulta necesario combinarlas.

TL;DR. El paralelismo de datos replicado permite aumentar el throughput de entrenamiento cuando una réplica cabe en una GPU. El paralelismo de datos totalmente sharded particiona el estado del modelo, pero añade all-gathers de parámetros y reduce-scatters de gradientes. El paralelismo tensorial, de pipeline, de contexto y de expertos divide, respectivamente, las operaciones de las capas, la profundidad, la secuencia y las capas mixture-of-experts (MoE). Combínalos solo después de identificar la restricción dominante de memoria o comunicación.

Esta guía presupone que tienes soltura con la backpropagation, las capas Transformer y un training loop estándar de PyTorch.

Empieza con dos presupuestos de memoria

El entrenamiento y la inferencia no tienen la misma huella de memoria.

training peak ≈ parameters
              + gradients
              + optimizer state
              + saved activations
              + temporary buffers
              + communication buffers
              + allocator headroom

inference peak ≈ resident weights
               + key/value (KV) cache
               + runtime workspace
               + communication buffers
               + allocator headroom

Un modelo de 70.000 millones de parámetros tiene un límite inferior decimal de 140 GB solo para los pesos en BF16. Esa cifra dice poco sobre el entrenamiento, donde los gradientes, el estado del optimizador, los pesos maestros y las activaciones pueden dominar el consumo. Tampoco permite dimensionar el serving, donde importan la política de caché, la longitud de la secuencia, la concurrencia del batch y la cuantización.

Perfila la arquitectura exacta, la precisión, la longitud de secuencia, el micro-batch, el optimizador, la política de checkpointing y el runtime. Registra la memoria máxima asignada y reservada, los tokens por segundo, el tiempo en kernels y el tiempo expuesto en collectives.

El escalado distribuido empieza por un cuello de botella medido

Cada dimensión de paralelismo implica un compromiso

Para cada estrategia, pregunta qué dimensión del tensor se divide, qué estado se replica y qué collective entra en el critical path.

Estrategia	Divide	Alivio principal	Comunicación introducida
Paralelismo de datos replicado	batch	throughput de entrenamiento	all-reduce de gradientes
Paralelismo de datos totalmente sharded	parámetros, gradientes y estado del optimizador dentro de un grupo data-parallel (DP)	memoria del estado del modelo	all-gather de parámetros, reduce-scatter de gradientes
Paralelismo tensorial	dimensiones de matrices o de atención dentro de las capas	pesos y activaciones de las capas	collectives dentro de los bloques Transformer
Paralelismo de pipeline	grupos de capas	profundidad del modelo y estado por stage	activaciones punto a punto más bubbles de scheduling
Paralelismo de contexto	dimensión de secuencia	memoria de activaciones de secuencias largas	intercambio de key/value o de atención entre el grupo de secuencia
Paralelismo de expertos	expertos MoE y tokens enrutados	capacidad de expertos por rank	dispatch y combine de tokens, normalmente all-to-all

El alivio de memoria no es un multiplicador fijo. Depende del grado de sharding, de lo que permanezca replicado, del estado transitorio sin sharding, de la política de activaciones, del padding, del desequilibrio y de los buffers.

Paralelismo de datos replicado: throughput sin aumentar la capacidad

Entrenamiento con paralelismo de datos replicado

El paralelismo de datos replicado, normalmente llamado distributed data parallel, mantiene una réplica completa del entrenamiento en cada rank. La documentación de DistributedDataParallel de PyTorch describe este modelo de réplica y sincronización de gradientes. Cada rank procesa un micro-batch diferente y los gradientes se sincronizan antes del paso del optimizador.

Úsalo cuando todo el estado del entrenamiento quepa con un margen de seguridad y el batch global pueda crecer, o cuando sea posible ajustar la acumulación de gradientes. Sus principales ventajas son una semántica sencilla y una implementación madura que solapa la computación backward con la reducción de gradientes agrupada en buckets.

Añadir ranks puede perjudicar el rendimiento cuando el batch local se vuelve demasiado pequeño. También puede hacerlo cuando la red no consigue ocultar el all-reduce o se producen esperas al entregar los datos de entrada. El batch de optimización deseado quizá no escale.

Paralelismo de datos totalmente sharded: memoria de estado a cambio de collectives

Ejecución de paralelismo de datos totalmente sharded

El paralelismo de datos totalmente sharded almacena los shards de parámetros, gradientes y estado del optimizador distribuidos entre un grupo. El artículo de ZeRO describe este patrón de sharding del estado, y la API FSDP2 de PyTorch lo implementa. Los parámetros de una capa se reúnen mediante all-gather para la computación y pueden volver a particionarse después. Los gradientes se devuelven a sus propietarios mediante reduce-scatter.

La documentación actual de PyTorch distingue la API fully_shard en la versión 2 del paralelismo de datos totalmente sharded (FSDP2) del wrapper FullyShardedDataParallel anterior. FSDP2 agrupa la comunicación según los módulos en los que se aplica fully_shard y recomienda aplicarlo de abajo arriba para que los grupos de capas puedan solapar comunicación y computación.

from torch.distributed.fsdp import fully_shard
from torch.optim import AdamW

# Apply bottom-up: each block becomes a communication group.
for block in model.transformer.blocks:
    fully_shard(block)

# Shard remaining root parameters such as embeddings and output projection.
fully_shard(model)

# Construct the optimizer after parameters have become sharded distributed tensors (DTensors).
optimizer = AdamW(model.parameters(), lr=learning_rate)

Este es un esquema estructural, no un launcher completo. Device meshes, mixed precision, checkpointing, inicialización, estado del optimizador y checkpoints distribuidos deben ser coherentes con el training stack.

El sharding resulta atractivo cuando el estado del modelo es la restricción dominante y la computación de las capas puede ocultar suficiente tráfico de collectives. Puede ser un mal intercambio para modelos pequeños, enlaces lentos, capas diminutas o layouts cuyo grupo de shards atraviese un límite de topología inadecuado.

Paralelismo tensorial: particionar las operaciones de las capas

Partición tensorial de una capa

El paralelismo tensorial particiona el álgebra lineal dentro de una capa. Algunos ejemplos son las proyecciones column-parallel y row-parallel. La guía de paralelismo de NVIDIA documenta esta división a nivel de capa. Los resultados parciales requieren collectives dentro de los bloques Transformer, por lo que la latencia y el ancho de banda importan repetidamente durante los pases forward y backward.

Úsalo cuando una capa o sus activaciones no quepan, o cuando las matmuls sean lo bastante grandes como para que los kernels particionados superen a un único rank. Asigna el grupo tensor-parallel al dominio de comunicación más rápido disponible y mide después. Un grado elevado de tensor parallelism puede reducir cada matriz local hasta que disminuya la eficiencia de los kernels, mientras aumenta el overhead de los collectives.

El sequence parallelism suele combinarse con el tensor parallelism para evitar replicar parte del trabajo sobre las activaciones. Es distinto del context parallelism sobre la secuencia de entrada completa del modelo.

Paralelismo de pipeline: particionar la profundidad y el tiempo de scheduling

Stages de pipeline y micro-batches

El paralelismo de pipeline coloca distintos grupos de capas en diferentes stages y envía las activaciones entre ellos. Los micro-batches mantienen los stages trabajando de forma concurrente. El artículo de GPipe utiliza este schedule para redes neuronales gigantes.

Reduce el estado del modelo por stage y puede disminuir el volumen de comunicación que cruza un límite más lento en comparación con los collectives tensoriales por capa. Sus costes son los bubbles, las transferencias de activaciones, el desequilibrio entre stages, un scheduling más complejo y una recuperación y un checkpointing más difíciles.

Para un schedule de estilo GPipe sencillo y equilibrado, con p stages y m micro-batches, la fracción idealizada del bubble forward es aproximadamente:

(p - 1) / (m + p - 1)

Los schedules reales pueden usar variantes one-forward/one-backward, interleaving o zero-bubble, y un coste desigual de las capas puede dominar la fórmula. Elige los límites de los stages a partir del tiempo y la memoria medidos, no del mismo número de capas en cada uno.

Paralelismo de contexto: particionar las activaciones de secuencias largas

Intercambio de atención en context parallelism

El context parallelism distribuye la dimensión de secuencia. La documentación de context parallelism de NVIDIA describe la división de la secuencia y el intercambio de key/value necesario para la atención. Cada rank posee un shard de la secuencia, mientras que la atención intercambia la información necesaria para preservar la semántica de contexto completo. Las implementaciones pueden usar anillos punto a punto, all-gather, all-to-all o combinaciones jerárquicas.

Reduce la memoria de activaciones durante el entrenamiento con contextos largos, pero replica los pesos dentro del grupo de contexto e introduce comunicación de atención. El beneficio depende del tipo de atención, la máscara causal, la longitud de secuencia, la recomputación y la forma en que los grupos de contexto se combinan con los grupos tensor-parallel y data-parallel.

No lo selecciones a partir de un umbral universal de 8K, 32K o 100K. Perfila la memoria de activaciones y la comunicación de atención para la arquitectura real.

Paralelismo de expertos: solo para una arquitectura MoE

Enrutamiento de tokens con expert parallelism

El expert parallelism distribuye los expertos de las capas mixture-of-experts. La guía de paralelismo de NVIDIA documenta esta colocación de expertos y su combinación con otras dimensiones de paralelismo. El router envía las representaciones de los tokens a los expertos seleccionados y combina sus resultados. Solo los expertos seleccionados calculan para cada token, pero el conjunto de pesos de los expertos sigue necesitando almacenamiento y una estrategia de colocación en serving.

El expert parallelism no es un interruptor de optimización para un modelo denso. Forma parte de una arquitectura MoE. Sus aspectos críticos incluyen el balanceo de carga, los límites de capacidad, el all-to-all de tokens, los tokens descartados o rellenados con padding, las pérdidas auxiliares y el desequilibrio ante fallos. Monitoriza los tokens por experto, la entropía del routing, el desbordamiento de capacidad, el tiempo de comunicación y la calidad según la ruta.

Compón el layout a partir de la topología

Los sistemas de entrenamiento de modelos densos grandes sin un grupo expert-parallel suelen usar el producto de los tamaños de los grupos data-parallel (DP), tensor-parallel (TP), pipeline-parallel (PP) y context-parallel (CP):

world size = DP × TP × PP × CP

Cuando expert parallelism (EP) es un grupo independiente, la guía de paralelismo de NVIDIA calcula el total así:

total GPUs = TP × PP × CP × EP × DP

Usa la mesh compatible con el framework en lugar de multiplicar configuraciones no compatibles.

Construye el layout en este orden:

Dibuja los dominios de comunicación: enlaces GPU-GPU, switches, límites de non-uniform memory access (NUMA), fabric del nodo, oversubscription y ruta de almacenamiento.
Coloca los collectives frecuentes y sensibles a la latencia, normalmente TP, en el dominio adecuado más rápido.
Elige los grupos FSDP o DP replicado a partir de la capacidad y el ancho de banda restantes.
Añade PP cuando la colocación por profundidad o el tráfico entre dominios lo justifiquen, equilibrando el tiempo y la memoria medidos de cada stage.
Añade CP solo por la restricción de secuencia, y expert parallelism (EP) únicamente por la topología de expertos del modelo.
Confirma la divisibilidad de cabezas, dimensiones ocultas, capas, expertos, batch y secuencia para la mesh candidata.
Haz benchmark de varias meshes válidas. Las heurísticas conscientes de la topología seleccionan candidatos, no ganadores.

Dos clusters con el mismo número de GPUs pueden preferir layouts distintos porque difieren el ancho de banda de los enlaces, la jerarquía de switches, la conexión de las CPUs y la congestión de red.

El entrenamiento y el serving requieren decisiones separadas

La inferencia normalmente no utiliza gradientes ni estado del optimizador, por lo que los layouts de entrenamiento de estilo FSDP no se transfieren automáticamente.

Para serving, pregunta:

¿Cabe una réplica con los pesos, el KV cache, el workspace y la concurrencia objetivo?
¿Se obtiene más throughput con más réplicas independientes o con el sharding de una sola réplica?
¿El TP reduce suficientemente la presión sobre los pesos y la caché de cada rank como para compensar la comunicación por capa?
¿El PP está implementado de forma eficiente para el modelo y el scheduler de peticiones?
¿Cómo someten prefill y decode a la computación, el ancho de banda de memoria y el interconnect?
¿Qué ocurre con la latencia de cola cuando las peticiones tienen longitudes diferentes de prompt y de salida?

Haz benchmark del servidor completo con el scheduler, la cuantización, la distribución del contexto, la política de batching y la forma del tráfico. Los tokens por segundo del entrenamiento no permiten predecir el tiempo hasta el primer token ni la latencia entre tokens en serving.

Mide un layout de escalado con rigor

Para cada candidato, registra:

identidad del modelo, el código, el runtime, los kernels y la topología
batch global y local, distribución de secuencias y número de tokens
memoria máxima por categoría, cuando esté disponible
tokens útiles por segundo y utilización de operaciones de coma flotante del modelo (FLOP), cuando se calculen de forma coherente
tiempo expuesto en operaciones all-reduce, all-gather, reduce-scatter, all-to-all y punto a punto
esperas de entrada, tiempo de checkpoint, comportamiento tras reinicios y distribución de stragglers
loss de entrenamiento o paridad de las salidas del serving con respecto al baseline

Compara deliberadamente el strong scaling y el weak scaling. El strong scaling mantiene constante el trabajo total a medida que aumenta el número de ranks. El weak scaling mantiene constante el trabajo por rank, por lo que el trabajo total crece con el número de ranks. Un porcentaje etiquetado como «eficiencia de escalado» carece de significado sin conocer ese denominador y el baseline.

Conclusión

El paralelismo es un mapeo entre un cuello de botella medido, una dimensión del tensor y un patrón de comunicación. La replicación, el sharding, la partición de capas, el staging, la partición de la secuencia y el routing de expertos alivian restricciones diferentes y crean modos de fallo distintos.

Haz inventario de la carga de trabajo, dibuja la topología, genera meshes válidas y perfílalas. El layout ganador es el que cabe con margen y minimiza la comunicación expuesta para el trabajo que realmente ejecutas.

Referencias

Documentación de PyTorch FSDP2 fully_shard - Sharding de parámetros a nivel de módulo, all-gathers, reduce-scatters y comportamiento de DTensor.
Documentación de PyTorch DistributedDataParallel - Entrenamiento de modelos replicados con gradientes sincronizados.
Estrategias de paralelismo de NVIDIA Megatron Core - Dimensiones de estrategia de datos, tensorial, de pipeline, de contexto, de expertos y totalmente sharded.
Ultra-Scale Playbook de Hugging Face - Visión general interactiva del paralelismo de modelos a gran escala.
Artículo de DeepSpeed sobre ZeRO - Sharding del estado del optimizador, los gradientes y los parámetros para reducir el uso de memoria.
Artículo sobre GPipe - Entrenamiento de pipeline con micro-batches y ejecución por stages.
Documentación de paralelismo de NERSC - Definiciones de strong scaling y weak scaling basadas en trabajo total o por rank constante.