Fully Sharded Data Parallel (FSDP) es un paradigma avanzado de entrenamiento de modelos de Machine Learning distribuido, diseñado para escalar el entrenamiento de modelos masivos que no caben en la memoria de una única unidad de procesamiento gráfico (GPU). A diferencia del Data Parallel tradicional, que replica el modelo completo en cada dispositivo y fragmenta solo los datos de entrada, FSDP fragmenta (shards) no solo los datos de entrada, sino también los parámetros del modelo, los gradientes y los estados del optimizador (como el momento o la varianza en Adam) a través de los dispositivos. Cada GPU solo almacena una porción del modelo, reduciendo drásticamente el uso de memoria por dispositivo. Durante el forward y backward pass, las porciones del modelo se recopilan dinámicamente (all-gather) en los dispositivos que las necesitan para realizar los cálculos, y luego se descartan, liberando memoria.
FSDP ha sido implementado y popularizado por Meta Platforms, siendo una característica clave en su biblioteca PyTorch. PyTorch FSDP es una implementación robusta que permite a los desarrolladores entrenar modelos de lenguaje grandes (LLMs) y otros modelos de aprendizaje profundo con miles de millones de parámetros de manera eficiente en clústeres de GPUs. Herramientas como Hugging Face Accelerate y DeepSpeed también ofrecen funcionalidades similares o integraciones con FSDP, facilitando su adopción en la comunidad de investigación y desarrollo de IA. Su uso es fundamental en el entrenamiento de modelos como Llama de Meta, donde la escala del modelo excede con creces la capacidad de memoria de un solo nodo o GPU.
Para un Arquitecto de Sistemas, FSDP es crucial porque permite superar las limitaciones de memoria de hardware, habilitando el entrenamiento de modelos de IA que antes eran inviables. La decisión de adoptar FSDP implica un trade-off: si bien reduce significativamente el consumo de memoria por GPU y permite escalar a modelos más grandes, introduce una sobrecarga de comunicación de red sustancial debido a las operaciones de 'all-gather' y 'reduce-scatter' entre dispositivos. Un arquitecto debe evaluar cuidadosamente la topología de red, el ancho de banda y la latencia del clúster de GPUs. La elección entre FSDP, Pipeline Parallelism o Tensor Parallelism dependerá de la escala del modelo, el número de GPUs disponibles, las características de la interconexión de red y los objetivos de rendimiento. FSDP es ideal cuando la memoria es el cuello de botella principal y se dispone de una red de alta velocidad entre los nodos.