Tensor Parallelism vs Pipeline Parallelism en LLMs: Cuándo Usar Cada Uno con PyTorch y vLLM

1. El Reto del Particionado en Modelos de Escala Frontera

Cuando un modelo de lenguaje excede la memoria física de una tarjeta gráfica (por ejemplo, Llama-3.1-70B en precisión bfloat16 requiere 140 GB solo para sus parámetros estáticos, sin contar KV Cache ni activaciones), la única alternativa es distribuir la red neuronal a lo largo de un clúster de procesadores aceleradores. Sin embargo, dividir una red neuronal profunda no es una tarea trivial: la forma en que se particionan los tensores y las capas determina drásticamente la latencia de respuesta, el uso de memoria y la saturación de los canales de comunicación de red.

2. Anatomía de Tensor Parallelism (TP)

Tensor Parallelism (formalizado por NVIDIA en la arquitectura Megatron-LM) no divide el modelo por capas consecutivas, sino que disecciona cada matriz interna dentro de cada bloque Transformer a través de las GPUs disponibles de forma coordinada.

En una capa Transformer clásica existen dos componentes esenciales:

1. Multi-Head Attention (MHA):
   • Las matrices de proyección de Query ($W_q$), Key ($W_k$) y Value ($W_v$) se dividen por columnas (Column Parallel). Cada GPU calcula las cabezas de atención correspondientes a su fragmento sin necesidad de comunicarse con las demás.
   • La matriz de proyección de salida ($W_o$) se divide por filas (Row Parallel). Para reconstruir la salida combinada, se ejecuta una sincronización colectiva obligatoria mediante una operación All-Reduce de NCCL.
2. Multi-Layer Perceptron (MLP / Feed-Forward):
   • La matriz de proyección ascendente (Gate/Up) se divide por columnas.
   • La matriz de proyección descendente (Down) se divide por filas, seguida de una segunda operación All-Reduce.

Esto significa que en un modelo con 80 capas Transformer, Tensor Parallelism ejecuta exactamente 160 sincronizaciones All-Reduce por cada token generado. Si la interconexión entre GPUs ofrece un ancho de banda descomunal (como los 900 GB/s de NVLink 4), el coste de sincronización es despreciable (~10 microsegundos). Pero si se ejecuta a través de ranuras PCIe Gen4 (64 GB/s) o redes TCP/IP, el sistema se colapsa de inmediato.

3. Anatomía de Pipeline Parallelism (PP)

Pipeline Parallelism adopta un esquema modular y desacoplado: divide la secuencia total de capas del modelo en bloques contiguos y asigna cada bloque a una GPU diferente conectada en cadena.

Por ejemplo, en un clúster de 4 GPUs ejecutando un modelo de 80 capas:

• GPU 0 (Etapa 1): Capa de embedding inicial + Capas 1 a 20.
• GPU 1 (Etapa 2): Capas 21 a 40.
• GPU 2 (Etapa 3): Capas 41 a 60.
• GPU 3 (Etapa 4): Capas 61 a 80 + Cabezal de proyección final (LM Head).

La ventaja primordial de PP es que la comunicación inter-GPU es estrictamente punto a punto (Point-to-Point): la GPU 0 simplemente transfiere el tensor de activación resultante de la capa 20 a la GPU 1. No se requiere ninguna operación All-Reduce masiva en cada subcapa.

Sin embargo, introduce el problema de la Burbuja de Pipeline (Pipeline Bubble): mientras la GPU 0 procesa el inicio de una secuencia, las GPUs 1, 2 y 3 permanecen inactivas esperando sus datos. Del mismo modo, una vez que la activación avanza a las últimas etapas, las primeras GPUs quedan ociosas. En inferencia secuencial de un solo usuario, la utilización de cómputo del clúster puede caer por debajo del 30%.

4. Tabla Comparativa de Arquitectura

5. Implementación en vLLM: Guía Práctica

vLLM soporta ambas técnicas con configuración directa mediante flags de línea de comandos:

Caso 1: Tensor Parallelism en 4 GPUs dentro del mismo servidor (Ideal)

vllm serve meta-llama/Meta-Llama-3.1-70B-Instruct \
  --tensor-parallel-size 4 \
  --pipeline-parallel-size 1 \
  --gpu-memory-utilization 0.90 \
  --port 8000

Caso 2: Servidor con 8 GPUs sin NVLink completo (Híbrido TP + PP)

# Configura 4 GPUs en TP dentro de cada par, y 2 etapas de Pipeline entre pares
vllm serve meta-llama/Meta-Llama-3.1-70B-Instruct \
  --tensor-parallel-size 4 \
  --pipeline-parallel-size 2 \
  --gpu-memory-utilization 0.92 \
  --port 8000

6. Diagnóstico de Problemas y Cuellos de Botella de Comunicación

1. El error NCCL error: unhandled system error / connection reset

Suele ocurrir cuando los adaptadores de red o los controladores de GPU tienen desajustes de versión con la biblioteca NCCL. En clústeres multi-GPU, asegúrate de configurar las variables de entorno de depuración:

export NCCL_DEBUG=INFO
export NCCL_IB_DISABLE=0 # Mantener habilitado InfiniBand si está disponible
export NCCL_P2P_DISABLE=0 # Habilitar comunicación directa Peer-to-Peer

2. Desbalance de carga en Pipeline Parallelism (Micro-batch Starvation)

Si la carga de secuencias no se divide en suficientes micro-batches (e.g. menos de 4 micro-batches por paso), la GPU de la última etapa pasará más del 50% del tiempo de ejecución esperando que las etapas iniciales completen su trabajo. Para mitigar esto en frameworks como Megatron-LM o vLLM, asegúrate de utilizar planificadores de pipeline interleaved tipo 1F1B (One-Forward-One-Backward) donde cada procesador alterna rápidamente entre cálculo forward y backward de distintas sub-etapas.

3. Picos de VRAM asimétricos entre GPUs en Tensor Parallelism

Aunque Tensor Parallelism reparte los pesos de forma simétrica entre las GPUs, la memoria consumida por las cabezas de embedding y la capa de salida LM Head puede generar un desbalance de hasta 4 GB adicionales en la GPU 0 si no se implementa tensor_parallel_split_embeddings=True. En vLLM esto se gestiona internamente, pero en scripts personalizados de PyTorch debes asegurar que la matriz de vocabulario de 128.000 tokens también se divida entre las tarjetas.

7. Configuración de Producción con Docker Compose Multi-GPU

A continuación se muestra un archivo docker-compose.yml listo para desplegar un cluster de inferencia con soporte de reservas de GPUs NVIDIA y optimización de memoria:

version: '3.8'

services:
  vllm-cluster:
    image: vllm/vllm-openai:latest
    container_name: vllm-prod
    ipc: host
    environment:
      - CUDA_VISIBLE_DEVICES=0,1,2,3
      - NCCL_DEBUG=WARN
      - PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
    ports:
      - "8000:8000"
    volumes:
      - ~/.cache/huggingface:/root/.cache/huggingface
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 4
              capabilities: [gpu]
    command: >
      --model meta-llama/Meta-Llama-3.1-70B-Instruct
      --tensor-parallel-size 4
      --max-model-len 8192
      --gpu-memory-utilization 0.92
      --port 8000

8. Checklist para Arquitectura de Sistemas Multi-GPU

• Comprueba con nvidia-smi topo -m la matriz de conectividad de tu servidor. Si las conexiones entre tarjetas indican NV# (NVLink), utiliza siempre Tensor Parallelism. Si indican SYS o PHB (PCIe Host Bridge), evita TP de más de 2 GPUs.
• Para entrenamiento y fine-tuning a gran escala, combina siempre TP dentro del nodo con Pipeline Parallelism y Zero Redundancy Optimizer (ZeRO) entre nodos remotos.

8. Optimización de Comunicaciones Colectivas NCCL sobre InfiniBand y RoCE

En clústeres multi-nodo donde se combinan Tensor Parallelism y Pipeline Parallelism, la red de interconexión es el componente más frágil. Una mala configuración de los buffers de NCCL puede degradar el rendimiento hasta en un 80%:

# Variables de entorno críticas para servidores con interfaces Mellanox ConnectX-6 / ConnectX-7
export NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3
export NCCL_IB_GID_INDEX=3
export NCCL_IB_TC=106
export NCCL_IB_SL=0
export NCCL_NET_GDR_LEVEL=5 # Habilitar GPUDirect RDMA directo entre VRAM y tarjeta de red
export NCCL_BUFFSIZE=8388608 # Aumentar el tamaño del buffer a 8 MB para saturar enlaces de 400 Gbps

La directiva NCCL_NET_GDR_LEVEL=5 garantiza que las tarjetas gráficas puedan transferir paquetes directamente a las tarjetas de red de alta velocidad sin hacer copias intermedias en la memoria RAM del sistema ni pasar por los núcleos de la CPU.

9. Planificación de Capacidad: Matriz de Asignación de GPUs según Modelo