de una serie sobre IA distribuida en múltiples GPU:
Introducción
En la publicación anterior, vimos cómo el paralelismo de datos distribuidos (DDP) acelera el entrenamiento al dividir lotes entre GPU. DDP resuelve el problema del rendimiento, pero introduce un nuevo desafío: la redundancia de memoria.
En Vanilla DDP, cada GPU contiene una copia completa de los parámetros del modelo, los gradientes y los estados del optimizador. Para modelos grandes como GPT-3 (parámetros 175B), esta redundancia se convierte en un gran desperdicio de la preciosa VRAM.
ZeRO (Zero Redundancy Optimizer) resuelve este problema. Hay tres niveles:
Estados del optimizador de particiones ZeRO-1 solo estados del optimizador de particiones ZeRO-2 + gradientes Estados del optimizador de particiones ZeRO-3 + gradientes + parámetros del modelo
ZeRO no es una técnica de paralelismo porque todas las GPU siguen ejecutando los mismos pases hacia adelante y hacia atrás. Es una estrategia de optimización de la memoria que elimina la redundancia entre GPU, lo que le permite entrenar modelos más grandes en el mismo hardware.
El problema de la memoria en DDP
Analicemos qué es lo que realmente consume memoria durante el entrenamiento. Para un modelo con parámetros:
Parámetros del modelo: valores (los pesos de su red neuronal) Gradientes: valores (un gradiente por parámetro) Estados del optimizador (Adam): valores (primer momento y segundo momento para cada parámetro) Activaciones: salidas intermedias almacenadas durante el paso hacia adelante para usar en el paso hacia atrás
Los primeros tres escalan con el tamaño del modelo y son redundantes en todas las GPU en DDP. Las activaciones escalan con el tamaño del lote, la longitud de la secuencia y el número de neuronas, y son únicas por GPU, ya que cada GPU procesa datos diferentes. ZeRO no toca la memoria de activación.
Calculemos el uso de memoria para un modelo de 7B parámetros usando Adam y FP32:
Parámetros: 7 mil millones * 4 bytes = 28 GB Degradados: 7 mil millones * 4 bytes = 28 GB Estados del optimizador: 7 mil millones * 2 * 4 bytes = 56 GB Memoria por GPU en DDP: 112 GB
Las activaciones añaden una cantidad significativa de memoria además de esto, pero como son únicas por GPU, ZeRO no puede particionarlas. Técnicas como el control de activación pueden ayudar, descarta algunas activaciones y luego las recalcula según sea necesario durante el paso hacia atrás. Pero eso está fuera del alcance de este artículo.
Entendamos cómo funciona ZeRO implementándolo desde cero, comenzando con ZeRO-1 y avanzando hasta ZeRO-3.
ZeRO-1: partición del estado del optimizador
En ZeRO-1, solo se particionan los estados del optimizador. Cada GPU:
Aún conserva los gradientes y parámetros completos del modelo. Almacena solo 1/N de los estados del optimizador (N = número de GPU) Actualiza solo el 1/N correspondiente de los parámetros.
Esta es la secuencia de acciones tomadas durante el entrenamiento:
Paso hacia adelante: cada GPU procesa su propio microlote Paso hacia atrás: calcula gradientes y reduce todos los gradientes: cada GPU obtiene todos los gradientes Paso del optimizador: cada GPU actualiza su partición de parámetros y reúne todos los parámetros: sincroniza el modelo actualizado entre las GPU
Aquí hay una implementación simplificada:
import torch import torch.distributed as dist class ZeRO_1: def __init__(self, model, optimizador_cls): self.model = model self.rank = dist.get_rank() self.world_size = dist.get_world_size() self.param_shards = list() # cada rango contiene solo su fragmento de los estados del optimizador self.param_metadata = list() # metadatos para reconstruir fragmentos para el parámetro en self.model.parameters(): original_shape = param.data.shape flat = param.data.view(-1) numel = flat.numel() resto = numel % self.world_size pad_size = (self.world_size – resto) % self.world_size acolchado_numel = numel + pad_size shard_size = acolchado_numel // self.world_size shard_start = self.rank * shard_size shard_end = shard_start + shard_size self.param_metadata.append( { "original_shape": original_shape, "numel": numel, "padded_numel": padded_numel, "shard_size": shard_size, "shard_start": shard_start, "shard_end": shard_end, }) si pad_size > 0: flat_padded = torch.cat([flat, flat.new_zeros(pad_size)]) else: flat_padded = flat shard = flat_padded[shard_start:shard_end].clone() shard.requires_grad_(True) self.param_shards.append(shard) self.optimizer = optimizador_cls(self.param_shards) def Training_step(self, entradas, objetivos, loss_fn): salida = self.model(entradas) # pérdida directa = loss_fn(salida, objetivos) # calcular la pérdida loss.backward() # retroceso self._sync_gradients() # reducir todos los gradientes en las GPU self.optimizer.step() # actualizar el fragmento local de parámetros self._sync_params() # todos reúnen los parámetros del modelo # borrar gradientes para el siguiente paso para el parámetro en self.model.parameters(): param.grad = Ninguno def _sync_gradients(self): para idx, param in enumerate(self.model.parameters()): meta = self.param_metadata[idx] dist.all_reduce(param.grad, op=dist.ReduceOp.SUM) param.grad /= self.world_size self.param_shards[idx].grad = param.grad.view(-1)[meta["shard_start"]:meta["shard_end"]] def _sync_params(self): para idx, parámetro en enumerate(self.model.parameters()): meta = self.param_metadata[idx] full_flat = torch.empty(meta["padded_numel"], dispositivo=param.device, dtype=param.dtype) dist.all_gather_into_tensor(output_tensor=full_flat, input_tensor=self.param_shards[idx].data, ) reconstructed = full_flat[:meta["numel"]].view(meta["original_shape"]) param.data.copy_(reconstruido)
Tenga en cuenta que la reducción total sincroniza todos los gradientes, pero cada GPU solo usa los gradientes para su propia partición de parámetros, lo que está sobrecomunicando. ZeRO-2 soluciona este problema fragmentando también los gradientes.
En la práctica, nunca usaría ZeRO-1 ya que ZeRO-2 le brinda mejores ahorros de memoria esencialmente al mismo costo. Pero aun así vale la pena repasarlo con fines de aprendizaje.
Memoria con ZeRO-1, modelo 7B, 8 GPU:
Parámetros: 28 GB (completamente replicado) Degradados: 28 GB (completamente replicado) Estados del optimizador: 56 GB/8 = 7 GB Total por GPU: 63 GB (en lugar de GB)
ZeRO-2: partición de gradiente
ZeRO-2 divide tanto los estados como los gradientes del optimizador. Dado que cada GPU solo actualiza una partición de parámetros, solo necesita los gradientes correspondientes.
ZeRO-1 utiliza all-reduce, lo que le da a cada GPU todos los gradientes. ZeRO-2 reemplaza esto con reducción de dispersión, cada GPU recibe solo los gradientes que realmente necesita. Esto ahorra memoria y ancho de banda de comunicación.
Pasos de entrenamiento:
Paso hacia adelante: cada GPU procesa su propio microlote Paso hacia atrás: calcular gradientes reducir-gradientes de dispersión: cada GPU obtiene solo su partición Paso del optimizador: cada GPU actualiza su partición de parámetros todos los parámetros recopilados: sincroniza el modelo actualizado entre las GPU
La implementación es muy similar a ZeRO-1, pero el paso de sincronización de gradiente utiliza reducción-dispersión en lugar de reducción total:
Pero espera, si cada GPU calcula todos los gradientes durante la backprop, ¿cómo ahorra esto realmente VRAM? He aquí cómo:
A medida que los gradientes de los parámetros se calculan capa por capa, se reducen-dispersan inmediatamente y se libera la copia local (nuestra implementación simplificada no realiza esto). Durante la retropropiedad, solo necesita el gradiente de la siguiente activación neuronal para calcular el gradiente del parámetro actual, es decir, no necesita todo el gráfico de gradiente. De esa manera, puedes liberar memoria para gradientes a medida que avanzas hacia atrás, manteniendo solo la partición asignada para cada GPU.
Memoria con ZeRO-2, modelo 7B, 8 GPU:
Parámetros: 28 GB (completamente replicado) Degradados: 28 GB / 8 = 3,5 GB Estados del optimizador: 56 GB / 8 = 7 GB Total por GPU: 38,5 GB (antes 112 GB)
ZeRO-3: partición de parámetros
Estados, gradientes y parámetros del optimizador de particiones ZeRO-3. Cada GPU almacena solo 1/N del estado completo del modelo.
Durante los pasos hacia adelante y hacia atrás, cada capa necesita sus parámetros completos, pero cada GPU solo almacena una fracción. Por lo tanto, recopilamos todos los parámetros justo a tiempo, los usamos y luego los descartamos inmediatamente después.
Pasos de entrenamiento:
Pase hacia adelante: recopile todos los parámetros de la capa de todas las GPU Ejecute el paso hacia adelante de la capa usando las activaciones de la capa anterior como entrada Descarte los parámetros recopilados (conserve solo la partición local) Repita estos pasos hasta que todas las capas estén terminadas Pase hacia atrás (por capa, a la inversa): recopile todos los parámetros de la capa nuevamente Calcule los gradientes para la capa actual usando los gradientes de activación de la siguiente capa Reduzca y disperse los gradientes (cada GPU mantiene su fragmento) Descarte los parámetros recopilados (conserve solo la partición local) Repita estos pasos hasta que todas las capas estén terminadas. Cada GPU ejecuta un paso del optimizador en su partición. No es necesario realizar una recopilación final, ya que los parámetros se recopilan capa por capa durante el paso hacia adelante.
Aquí hay una implementación simplificada:
class ZeRO_3(ZeRO_2): """ ZeRO-3: estados del optimizador de fragmentos (etapa 1) + gradientes (etapa 2) + parámetros del modelo (etapa 3). En reposo, cada rango contiene solo param_shards[idx], una porción de 1/tamaño mundial de cada parámetro. Los parámetros completos se materializan temporalmente durante los pases hacia adelante y hacia atrás a través de all_gather, luego se liberan inmediatamente. """ def __init__(self, model, optimizador_cls): self.model = modelo self.rank = dist.get_rank() self.world_size = dist.get_world_size() self.param_metadata =[]lista_fragmentos =[]self._param_to_idx = {} para idx, parámetro en enumerate(self.model.parameters()): original_shape = param.data.shape flat = param.data.view(-1) numel = flat.numel() resto = numel % self.world_size pad_size = (self.world_size – resto) % self.world_size acolchado_numel = numel + pad_size shard_size = acolchado_numel // self.world_size shard_start = self.rank * shard_size shard_end = shard_start + shard_size self.param_metadata.append( { "original_shape": original_shape, "numel": numel, "padded_numel": acolchado_numel, "shard_size": shard_size, "shard_start": shard_start, "shard_end": shard_end, } ) si pad_size > 0: flat_padded = torch.cat([flat, flat.new_zeros(pad_size)]) else: flat_padded = flat shard = flat_padded[shard_start:shard_end].clone() shard_list.append(shard) # Reemplace el tensor completo solo con el fragmento de este rango. # Los datos param.data del modelo ahora apuntan a una pequeña porción; el # peso completo se reconstruirá según sea necesario durante el avance/retroceso. param.data = shard.detach() self._param_to_idx[param] = idx self.param_shards = [s.requires_grad_(True) para s en shard_list] self.optimizer = optimizador_cls(self.param_shards) self._register_hooks() def _gather_param(self, idx, dispositivo, dtype): """Reúna todo el tensor de parámetros completo para el parámetro `idx`.""" meta = self.param_metadata[idx] full_flat = torch.empty(meta["padded_numel"], dispositivo=dispositivo, dtype=dtype) dist.all_gather_into_tensor(output_tensor=full_flat, input_tensor=self.param_shards[idx].data, ) return full_flat[: meta["numel"]].view(meta["original_shape"]) def _gather_module_params(self, module): """Reúna los parámetros completos para cada parámetro que pertenece solo a este módulo (no los secundarios).""" for param in module.parameters(recurse=False): idx = self._param_to_idx[param] param.data = self._gather_param(idx, param.device, param.dtype) def _reshard_module_params(self, module): """Volver a dividir los parámetros en el fragmento local para cada parámetro directo de este módulo.""" for param in module.parameters(recurse=False): idx = self._param_to_idx[param] param.data = self.param_shards[idx].data def _register_hooks(self): self._hooks =[]para módulo en self.model.modules(): # Omitir módulos contenedores que no tienen parámetros directos si no list(module.parameters(recurse=False)): continuar # Adelante: recopilar -> ejecutar -> reshard h1 = module.register_forward_pre_hook( lambda mod, _inputs: self._gather_module_params(mod) ) h2 = module.register_forward_hook( lambda mod, _inputs, _output: self._reshard_module_params(mod) ) # Hacia atrás: recopilar antes del cálculo de graduación → reshard después de h3 = module.register_full_backward_pre_hook( lambda mod, _grad_output: self._gather_module_params(mod) ) h4 = module.register_full_backward_hook( lambda mod, _grad_input, _grad_output: self._reshard_module_params(mod) ) self._hooks.extend([h1, h2, h3, h4]) def Training_step(self, inputs, target, loss_fn): # Los ganchos manejan todas las recopilaciones/reestructuraciones alrededor de cada módulo automáticamente salida = self.model(inputs) loss = loss_fn(output, target) loss.backward() self._sync_gradients() # Cada rango actualiza solo su fragmento local self.optimizer.step() para param en self.model.parameters(): param.grad = Ninguno
Los parámetros de cada capa se recopilan justo antes de que sean necesarios y se liberan inmediatamente después. Esto mantiene la memoria máxima al mínimo a costa de una mayor comunicación. En la práctica, las implementaciones se superponen a la capa N+1 con el avance de la capa N para ocultar la latencia.
Memoria con ZeRO-3, modelo 7B, 8 GPU:
Parámetros: 28 GB / 8 = 3,5 GB Degradados: 28 GB / 8 = 3,5 GB Estados del optimizador: 56 GB / 8 = 7 GB Total por GPU: 14 GB (antes 112 GB)
Esto supone una reducción de 8 veces en el uso de memoria, que es exactamente lo que esperaríamos de la partición en 8 GPU.
Usando ZeRO en PyTorch
PyTorch se entrega con dos implementaciones de ZeRO-3: FSDP1 (más antigua, menos optimizada) y FSDP2 (más nueva, recomendada). Utilice siempre FSDP2.
FSDP (Fully Sharded Data Parallel) maneja automáticamente la recopilación de parámetros, la dispersión de gradientes, la superposición de comunicaciones y la administración de memoria:
de torch.distributed.fsdp importar modelo full_shard = Transformer() para capa en model.layers:fully_shard(capa)fully_shard(modelo)
Tienes que aplicar full_shard capa por capa y luego envolver todo el modelo.
Conclusión
ZeRO intercambia memoria por comunicación, por lo que no es un almuerzo gratis. En general, no vale la pena para modelos más pequeños (por ejemplo, BERT), pero cambia las reglas del juego para modelos más grandes.
¡Felicitaciones por llegar hasta el final! En esta publicación, aprendiste sobre:
El problema de redundancia de memoria en DDP estándar Cómo ZeRO particiona los estados, gradientes y parámetros del optimizador entre GPU Las tres etapas de ZeRO y sus compensaciones entre memoria y comunicación Cómo usar ZeRO-3 a través del FSDP de PyTorch
En el próximo artículo, exploraremos el paralelismo tensorial, una técnica de paralelismo de modelos que acelera el cálculo de una capa mediante la distribución del trabajo entre las GPU.
Referencias
ZeRO: Optimizaciones de memoria para el entrenamiento de billones de modelos de parámetros (artículo original) Tutorial de FSDP de PyTorch Referencia de API de FSDP El libro de estrategias de ultraescala de Hugging Face