¿Como funciona?
Como dije, cuando se entrena en varias GPU, cada proceso tiene copias exactas de los mismos datos cuando se entrena con DDP. Podemos optimizarlo implementando varias mejoras:
Estado del optimizador de fragmentos (ZeRO 1)
Cuando se entrena con DDP, cada proceso contiene una copia completa de los estados del optimizador. Con ZeRO1, fragmentamos estos estados del optimizador en todos los rangos de modo que cada rango contenga solo una parte de los estados del optimizador. Durante el paso hacia atrás, cada rango solo necesita recopilar los estados del optimizador relevantes para sus parámetros para realizar un paso de optimización. Esta reducción de la redundancia ayuda a conservar la memoria.
💡 En el caso de Adam, que mantiene parámetros de aproximadamente el doble del tamaño del modelo, dividir el estado del optimizador entre 8 rangos significa que cada rango almacena sólo una cuarta parte (2/8) del tamaño total del estado.
Degradados de fragmentos (ZeRO 2)
Fragmentamos los estados del optimizador. Ahora, modificaremos también el paso del optimizador para fragmentar los gradientes. Si un rango tiene estados de optimizador para una parte de los parámetros, entonces haremos lo siguiente:
- agregar todos los gradientes relevantes para los estados en los que se encuentra el rango
- calcular el paso de optimización
- enviar el paso de optimización para una parte de los parámetros a todos los demás rangos
Como habrás notado, ahora cada rango no necesita contener una réplica completa de los gradientes. Podemos enviar gradientes a un rango relevante tan pronto como estén disponibles. Por lo tanto, podemos reducir aún más el consumo máximo de memoria.
Parámetros del modelo de fragmentos (ZeRO 3)
Esto está a punto de ser épico.
¿Por qué necesitamos almacenar una copia completa del modelo en cada rango? Dividamos los parámetros del modelo entre todos los rangos. Luego, vamos a buscar los parámetros requeridos justo a tiempo durante el avance y el retroceso.
💡 En el caso de modelos grandes, estas optimizaciones pueden reducir drásticamente el consumo de memoria.
¿Cómo utilizar FSDP?
En realidad, bastante simple. Todo lo que necesitamos es envolver el modelo con FSDP:
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributed.fsdp import FullyShardedDataParallel as FSDPmodel = FSDP(model)
# it's critical to get parameters from the wrapped model
# as only a portion of them returned (sharded part)
optimizer = optim.Adam(model.parameters())
# consuct training as usual
train(model, optimizer)
También puede especificar la estrategia de fragmentación de FSDP. Por ejemplo, podemos seleccionar el SHARD_GRAD_OP estrategia para lograr un comportamiento similar al de ZeRO2. Puedes conocer otras estrategias aquí:
Además, puede empaquetar con submódulos FSDP. En el ejemplo anterior, solo se utiliza un módulo FSDP, lo que reducirá la eficiencia informática y la eficiencia de la memoria. La forma en que funciona es así, supongamos que su modelo contiene 100 capas lineales. Si utiliza FSDP (modelo), solo habrá una unidad FSDP que envuelva todo el modelo. En ese caso, allgather recopilaría los parámetros completos para las 100 capas lineales y, por lo tanto, no guardará memoria CUDA para la fragmentación de parámetros.
Puede ajustar submódulos explícitamente o definir una política de ajuste automático. Para obtener más información sobre FSDP, lea la guía de PyTorch: