English · Español

02 — Tensor parallel, pipeline parallel, sequence parallel, expert parallel¶

🇪🇸 La familia "el modelo está partido entre las GPUs". Tensor parallel parte dentro de una capa (la matriz de pesos se reparte por filas o por columnas). Pipeline parallel parte entre capas (capa 1 en GPU0, capa 2 en GPU1). Sequence parallel parte la dimensión de secuencia. Expert parallel parte expertos de MoE. Cada uno minimiza un coste distinto.

Mientras que la familia paralelismo de datos (data parallel) de theory/01-data-parallel-and-zero.md replica (o shardea-y-luego-junta) el modelo entre workers, la familia model-parallel parte el modelo en sí. Workers distintos computan partes distintas del mismo forward pass.

Cubrimos cuatro variantes: tensor parallel (TP), pipeline parallel (PP), sequence parallel, expert parallel. De estas, TP y PP son las que Borja necesita leer con soltura. Sequence parallel y expert parallel reciben cobertura conceptual.

Tensor parallel (TP) — estilo Megatron¶

Una única matriz de pesos se shardea entre \(N\) workers. Los dos splits naturales para una capa Linear(in_features=I, out_features=O):

Linear column-parallel¶

La dimensión de salida se parte: cada worker mantiene \(O/N\) columnas de salida de \(W\). Cómputo:

Input: la misma activación \(x \in \mathbb{R}^{B \times I}\) se replica en cada worker (así no hay comunicación antes de esta capa).
Cómputo local: el worker \(i\) computa \(y_i = x W_i \in \mathbb{R}^{B \times O/N}\) donde \(W_i\) es el shard de columnas del worker.
Output: cada worker tiene un shard de columnas distinto de \(y \in \mathbb{R}^{B \times O}\). Si la siguiente capa es row-parallel, no hay comunicación aquí — la siguiente capa espera este split.

Linear row-parallel¶

La dimensión de entrada se parte: cada worker mantiene \(I/N\) filas de entrada de \(W\). Cómputo:

Input: la activación \(x\) debe estar pre-partida — cada worker tiene \(x_i \in \mathbb{R}^{B \times I/N}\) (esto coincide con cómo se ve la salida de una capa column-parallel).
Cómputo local: el worker \(i\) computa \(y_i = x_i W_i \in \mathbb{R}^{B \times O}\) — dim de salida completa, pero solo una suma parcial.
Output: un all-reduce entre workers produce la \(y\) completa.

El patrón MLP¶

Un MLP transformer son dos linears con una activación en medio: Linear(d_model, d_ff) → GELU → Linear(d_ff, d_model). El patrón Megatron:

Capa 1: column-parallel — la salida es [B, d_ff/N] por worker.
GELU: elemento a elemento, sin comunicación.
Capa 2: row-parallel — la entrada es [B, d_ff/N] por worker, la salida es [B, d_model] completa tras un all-reduce.

Comunicación total: un all-reduce por MLP. Volumen: \(B \cdot d_{\text{model}} \cdot 4\) bytes (fp32) o \(\cdot 2\) (fp16). Para un transformer con \(L\) capas y un MLP por capa, eso son \(L\) all-reduces por forward pass (más otros \(L\) para el backward).

El patrón de atención¶

Las proyecciones Q, K, V son column-parallel (dim de salida partida: cada worker recibe n_heads / N heads). La proyección de salida Linear(d_model, d_model) es row-parallel. Mismo patrón que el MLP: 2 all-reduces por capa (forward + backward).

Shardear la tabla de embeddings¶

Específicamente para el grammar tutor: con ~600 formas (inglés + español), la tabla de embeddings es minúscula. Pero el patrón importa para la variante "creció en el futuro". El embedding nn.Embedding(vocab_size, d_model) shardea a lo largo de vocab_size:

El worker \(i\) tiene los embeddings para los token IDs \([i \cdot V/N, (i+1) \cdot V/N)\).
Un lookup: cada worker recoge su rebanada, luego all-reduce para combinar.

Para el grammar tutor con vocab=600, esto es cómicamente derrochador — el overhead del all-reduce empequeñece el lookup en sí. El lab lo hace de todas formas, como ejercicio docente, con la nota explícita "esto es desproporcionado para nuestra tarea; lo hacemos para ver la forma". Las cargas reales usan este patrón cuando vocab × \(d_{\text{model}}\) supera la memoria de una sola GPU (en algún momento alrededor de 600k tokens × \(d_{\text{model}} = 4096\) en fp32 ≈ 10 GB).

Pipeline parallel (PP)¶

Las capas se particionan en stages. El stage \(s\) vive en el worker \(s\). Un microbatch fluye worker \(0 \to 1 \to 2 \to \ldots \to S-1\) durante el forward, y luego de vuelta durante el backward.

El problema de la burbuja¶

Ingenuamente, mientras el worker 0 hace forward del microbatch 1, los workers 1..\(S-1\) están ociosos. Luego el worker 1 hace forward del microbatch 1 mientras el worker 0 empieza el forward del microbatch 2, etc. El "fill" y el "drain" del pipeline son tiempo ocioso.

Para \(S\) stages y \(M\) microbatches, la fracción de burbuja es:

\[\text{bubble} = \frac{S - 1}{M + S - 1}\]

Ejemplos:

\(S=4, M=4\): burbuja = 3/7 ≈ 43%. La mitad del cluster está ociosa la mitad del tiempo.
\(S=4, M=16\): burbuja = 3/19 ≈ 16%.
\(S=4, M=64\): burbuja = 3/67 ≈ 4%.

El número de microbatches debe superar masivamente el número de stages. Esta es la regla más violada de los rollouts PP en producción.

Scheduling 1F1B¶

El schedule ingenuo hace todos los forwards, luego todos los backwards. 1F1B (one-forward-one-backward) intercala: tan pronto como un microbatch termina forward en el último stage, empieza su backward. Esto reduce la memoria pico de activaciones (no tienes que mantener simultáneamente las activaciones de \(M\) microbatches) sin cambiar la burbuja.

1F1B intercalado¶

El schedule intercalado de Megatron: cada worker tiene múltiples stages no contiguos. El worker 0 tiene las capas 0–7 y 16–23; el worker 1 tiene 8–15 y 24–31. El forward pasa por el doble de stages, reduciendo a la mitad el cómputo de cada "stage", permitiendo más solape de pipeline. Reduce la burbuja a ~la mitad a costa de más comunicación.

El diagrama del timeline PP¶

Va en diagrams/. El lab lo genera. El esbozo mermaid:

tiempo → → →
W0: [F1 F2 F3 F4 .. .. .. .. B4 B3 B2 B1]   <-- ingenuo
W1: [.. F1 F2 F3 F4 .. .. B4 B3 B2 B1 ..]
W2: [.. .. F1 F2 F3 F4 B4 B3 B2 B1 .. ..]
W3: [.. .. .. F1 F2 F3 B3 B2 B1 .. .. ..]
       \________/         \______/
        burbuja de llenado burbuja de vaciado

con "F\(m\)" = forward del microbatch \(m\), "B\(m\)" = backward.

Sequence parallel¶

Las activaciones crecen como \(B \times L_{\text{seq}} \times d_{\text{model}}\). Para contextos largos \(L_{\text{seq}}\) se vuelve enorme. Sequence parallel shardea a lo largo de \(L_{\text{seq}}\) — workers distintos tienen posiciones distintas de la misma secuencia.

La pega: la atención no es local a la posición de secuencia. Para computar la atención en la posición \(i\), necesitas K/V en las posiciones \(1..i\) (causal). La atención sequence-parallel requiere comunicación inter-worker durante la propia atención.

Existen variantes: sequence parallel de Megatron (aumento modesto de comunicación), Ring Attention (los bloques KV rotan entre workers en un anillo) y DeepSpeed Ulysses (atención sequence-parallel con comunicación all-to-all). Todas chocan con un techo de contexto largo que escala más allá que TP solo.

Para el grammar tutor (contexto máximo ~32 tokens), sequence parallel es desproporcionado. Solo cobertura conceptual.

Expert parallel¶

Usado con Mixture-of-Experts (MoE). Una capa MoE tiene \(E\) expertos; cada token se enruta al top-\(k\) de ellos (típicamente \(k=2\)). Expert parallel asigna cada experto a un worker distinto.

Los tokens se enrutan vía comunicación all-to-all: cada worker manda sus tokens al worker que aloja el experto elegido, el experto computa, los resultados van de vuelta por all-to-all.

Expert parallel es el paralelismo para MoE; es por lo que MoE funciona a escala. La Fase 36 cubre MoE en profundidad. La Fase 35 lo menciona para vocabulario.

Resumen del paralelismo 3D¶

El entrenamiento de modelos grandes en producción combina:

DP entre nodos (el más externo) — replica el modelo (ya shardeado).
TP dentro de un nodo (el más interno) — usa el ancho de banda intra-nodo de NVLink de ~600 GB/s para los all-reduces conversadores de TP.
PP entre nodos cuando el escalado DP se satura — mueve activaciones de frontera de capa entre nodos (menos ancho de banda necesario que TP).

Una configuración común para un modelo 70B en un cluster de 8 nodos × 8 GPU: DP=2, TP=8 (intra-nodo), PP=4 (entre pares de nodos). La matemática es implacable — equivócate en cualquiera de los tres tamaños y el rendimiento colapsa.

Para la Fase 35: cubrimos las piezas. No implementamos la combinación 3D. Megatron-LM y DeepSpeed sí, y leer sus documentos de configuración (lab 03) es el aprendizaje.

¿Cuándo es correcto el model-parallel?¶

Situación	Recomendación
El modelo cabe en una GPU	No uses model parallel. Usa DDP o single-GPU.
El modelo apenas no cabe (5–20% por encima)	Prueba ZeRO-3 / FSDP primero. Más simple, más rápido de cablear.
Modelo 2–8× demasiado grande para una GPU	TP dentro de un solo nodo (NVLink intra-nodo).
Modelo demasiado grande para un nodo	TP + PP, con TP intra-nodo y PP entre nodos.
MoE con muchos expertos	Expert parallel, siempre.
Contexto largo (≫ 8k tokens)	Añade sequence parallel encima.

Para el grammar tutor de este currículo:

El modelo cabe en una calculadora. Realmente no necesita TP.
El lab 02 corre TP en 2 GPUs cloud como ejercicio docente: shardea la tabla de embeddings, observa el all-reduce, mide el speedup vs single-GPU (será un slowdown porque el modelo es demasiado pequeño — eso forma parte de la lección).

Lo que esta fase NO cubre¶

Escribir a mano un scheduler PP 1F1B intercalado. Existe torch.distributed.pipeline.sync de PyTorch; existe megatron/core/pipeline_parallel/schedules.py de Megatron. Leemos uno, no escribimos uno.
Implementaciones de Ring Attention / DeepSpeed Ulysses. Solo conceptual.
Implementación del routing de expertos MoE. Territorio de la Fase 36.
Compresión de comunicación (1-bit Adam, PowerSGD). Solo vocabulario.

Siguiente: theory/03-collectives-and-cost.md.