Sharding en LLM: cómo distribuir computaciones entre GPUs

El escalado de grandes modelos de lenguaje requiere más de un GPU o TPU. Una de las técnicas clave para esto es el sharding: la distribución de matrices y vectores entre múltiples aceleradores para que los cálculos se realicen de manera eficiente y coordinada.

Por qué las matrices no caben en la memoria

Al trabajar con redes neuronales verdaderamente grandes, las matrices de pesos y activaciones casi nunca caben completamente en la memoria de un solo GPU o TPU. El peso de un LLM moderno puede medirse en cientos de miles de millones de parámetros, cada uno requiriendo varios bytes de memoria. Un único acelerador simplemente no puede manejar esto. Por lo tanto, es necesario dividir —shardear— estas matrices entre varios chips.

Cómo funciona el sharding

Imaginemos una matriz A de tamaño [I, J]. Puede dividirse de varias maneras:

• Por filas (row-wise): el primer acelerador toma las filas 1–N, el segundo las filas N+1–2N y así sucesivamente
• Por columnas (column-wise): de manera similar, pero verticalmente
• Particionamiento en bloques: la matriz se divide en bloques rectangulares y se distribuye en una cuadrícula de chips
• Por diagonal y otros patrones: esquemas más complejos para operaciones específicas

La función de cada acelerador es realizar la operación en su parte de la matriz y luego sincronizar los resultados con los vecinos.

Cómo afecta la calidad del sharding

La elección correcta de la estrategia de sharding determina:

• Velocidad de entrenamiento: minimizar el tiempo de comunicación entre GPUs
• Eficiencia de memoria: llenado correcto de búferes en cada chip
• Ancho de banda: cantidad de datos que deben transmitirse entre aceleradores
• Localidad computacional: cuánto trabajo puede ejecutarse sin sincronización

Un sharding deficiente puede llevar a una situación donde los aceleradores pasan la mayor parte del tiempo esperándose mutuamente o desperdician el ancho de banda de las conexiones entre ellos. Un buen sharding intenta minimizar estas pérdidas y mantener todos los chips funcionando simultáneamente.

Qué significa esto

El sharding no es simplemente una optimización técnica, es la base para la economía del escalado. Una distribución correcta de los cálculos permite utilizar clusters de cientos y miles de GPUs de manera que la velocidad total se aproxime a la ideal y los recursos computacionales no se desperdicien en esperas de comunicación.