Paralelismo de datos totalmente fragmentado

Para acelerar el entrenamiento de modelos enormes con tamaños de lote más grandes, podemos usar un modelo de paralelismo de datos totalmente fragmentado. Este tipo de paradigma de paralelismo de datos permite ajustar más datos y modelos más grandes al fragmentar los estados del optimizador, los gradientes y los parámetros. Para leer más al respecto y sobre sus beneficios, consulta el blog de Fully Sharded Data Parallel. Hemos integrado la función de entrenamiento Fully Sharded Data Parallel (FSDP) más reciente de PyTorch. Todo lo que necesitas hacer es activarla a través de la configuración.

Cómo funciona sin configuración adicional

En su(s) máquina(s), simplemente ejecute:

accelerate config

y responder a las preguntas realizadas. Esto generará un archivo de configuración que se usará automáticamente para establecer correctamente las opciones predeterminadas al

accelerate launch my_script.py --args_to_my_script

Por ejemplo, así es como ejecutarías examples/nlp_example.py (desde la raíz del repositorio) con FSDP habilitado:

compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch_policy: BACKWARD_PRE
  fsdp_forward_prefetch: false
  fsdp_cpu_ram_efficient_loading: true
  fsdp_offload_params: false
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_transformer_layer_cls_to_wrap: BertLayer
  fsdp_use_orig_params: true
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

accelerate launch examples/nlp_example.py

Actualmente, Accelerate admite la siguiente configuración a través de la CLI:

fsdp_sharding_strategy: [1] FULL_SHARD (fragmenta los estados del optimizador, los gradientes y los parámetros), [2] SHARD_GRAD_OP (fragmenta los estados del optimizador y los gradientes), [3] NO_SHARD (DDP), [4] HYBRID_SHARD (fragmenta los estados del optimizador, los gradientes y los parámetros dentro de cada nodo, mientras que cada nodo tiene una copia completa), [5] HYBRID_SHARD_ZERO2 (fragmenta los estados del optimizador y los gradientes dentro de cada nodo, mientras que cada nodo tiene una copia completa). Para obtener más información, consulte la documentación oficial de PyTorch.

fsdp_offload_params : Decide si se deben descargar los parámetros y los gradientes a la CPU

fsdp_auto_wrap_policy: [1] TRANSFORMER_BASED_WRAP, [2] SIZE_BASED_WRAP, [3] NO_WRAP

fsdp_transformer_layer_cls_to_wrap: Solo aplicable para Transformers. Al usar fsdp_auto_wrap_policy=TRANSFORMER_BASED_WRAP, el usuario puede proporcionar una cadena separada por comas con los nombres de las clases de capas de transformer (distingue entre mayúsculas y minúsculas) que se deben envolver, por ejemplo, BertLayer, GPTJBlock, T5Block, BertLayer,BertEmbeddings,BertSelfOutput. Esto es importante porque los submódulos que comparten pesos (por ejemplo, las capas de embedding) no deben terminar en unidades FSDP envueltas diferentes. Al usar esta política, el envoltorio se realiza para cada bloque que contiene Multi-Head Attention seguido de un par de capas MLP. Las capas restantes, incluidas las embeddings compartidas, se envuelven de manera conveniente en la misma unidad FSDP más externa. Por lo tanto, úsela para modelos basados en transformer. Puede usar model._no_split_modules para los modelos de Transformer respondiendo yes a Do you want to use the model's_no_split_modulesto wrap. It will try to usemodel._no_split_modules` when possible.

fsdp_min_num_params: número mínimo de parámetros al usar fsdp_auto_wrap_policy=SIZE_BASED_WRAP.

fsdp_backward_prefetch_policy: [1] BACKWARD_PRE, [2] BACKWARD_POST, [3] NO_PREFETCH

fsdp_forward_prefetch: si es True, FSDP realiza explícitamente la preobtención del próximo all-gather mientras se ejecuta la pasada hacia adelante. Solo debe usarse con modelos de grafo estático, ya que la preobtención sigue el orden de ejecución de la primera iteración. Es decir, si el orden de los submódulos cambia dinámicamente durante la ejecución del modelo, no se debe habilitar esta función.

fsdp_state_dict_type: [1] FULL_STATE_DICT, [2] LOCAL_STATE_DICT, [3] SHARDED_STATE_DICT

fsdp_use_orig_params: Si es True, permite un requires_grad no uniforme durante la inicialización, lo que significa que se admiten parámetros congelados y entrenables intercalados. Esta configuración es útil en casos como el ajuste fino eficiente en parámetros, tal como se discutió en esta publicación. Esta opción también permite tener varios grupos de parámetros del optimizador. Debería ser True al crear un optimizador antes de preparar/envolver el modelo con FSDP.

fsdp_cpu_ram_efficient_loading: Solo aplicable a modelos de Transformers. Si es True, solo el primer proceso carga el checkpoint del modelo preentrenado, mientras que todos los demás procesos tienen pesos vacíos. Esto debe establecerse en False si experimenta errores al cargar el modelo preentrenado de Transformers mediante el método from_pretrained. Cuando esta configuración es True, fsdp_sync_module_states también debe ser True; de lo contrario, todos los procesos excepto el principal tendrían pesos aleatorios, lo que provocaría un comportamiento inesperado durante el entrenamiento. Para que esto funcione, asegúrese de que el grupo de procesos distribuidos esté inicializado antes de llamar al método from_pretrained de Transformers. Al usar la API de Trainer, el grupo de procesos distribuidos se inicializa al crear una instancia de la clase TrainingArguments.

fsdp_sync_module_states: Si es True, cada unidad FSDP envuelta individualmente difundirá los parámetros del módulo desde el rank 0.

Para un control adicional y más matizado, puedes especificar otros parámetros de FSDP a través de FullyShardedDataParallelPlugin. Al crear el objeto FullyShardedDataParallelPlugin, pásale los parámetros que no formaban parte de la configuración de accelerate o que deseas sobrescribir. Los parámetros de FSDP se seleccionarán en función del archivo de configuración de accelerate o de los argumentos del comando de lanzamiento, y los demás parámetros que pases directamente a través del objeto FullyShardedDataParallelPlugin los establecerán o sobrescribirán.

A continuación se muestra un ejemplo:

from accelerate import FullyShardedDataParallelPlugin
from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig

fsdp_plugin = FullyShardedDataParallelPlugin(
    state_dict_config=FullStateDictConfig(offload_to_cpu=False, rank0_only=False),
    optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=False, rank0_only=False),
)

accelerator = Accelerator(fsdp_plugin=fsdp_plugin)

Guardar y cargar

La nueva forma recomendada de crear puntos de control al usar modelos FSDP es utilizar SHARDED_STATE_DICT como StateDictType al configurar el archivo de configuración de accelerate. A continuación se muestra el fragmento de código para guardar usando la utilidad save_state de accelerate.

accelerator.save_state("ckpt")

Inspeccione la carpeta de checkpoint para ver el modelo y el optimizador como fragmentos por proceso:

ls ckpt
# optimizer_0  pytorch_model_0  random_states_0.pkl  random_states_1.pkl  scheduler.bin

cd ckpt

ls optimizer_0
# __0_0.distcp  __1_0.distcp

ls pytorch_model_0
# __0_0.distcp  __1_0.distcp

Para cargarlos de nuevo y reanudar el entrenamiento, utilice la utilidad load_state de accelerate

accelerator.load_state("ckpt")

Al usar save_pretrained de transformers, pasa state_dict=accelerator.get_state_dict(model) para guardar el state dict del modelo. A continuación se muestra un ejemplo:

  unwrapped_model.save_pretrained(
      args.output_dir,
      is_main_process=accelerator.is_main_process,
      save_function=accelerator.save,
+     state_dict=accelerator.get_state_dict(model),
)

State Dict

accelerator.get_state_dict llamará a la implementación subyacente de model.state_dict utilizando el gestor de contexto FullStateDictConfig(offload_to_cpu=True, rank0_only=True) para obtener el state dict solo para el rank 0 y se descargará a la CPU.

A continuación, puede pasar state al método save_pretrained. Existen varios modos para StateDictType y FullStateDictConfig que puede utilizar para controlar el comportamiento de state_dict. Para obtener más información, consulte la documentación de PyTorch.

Si eliges usar StateDictType.SHARDED_STATE_DICT, los pesos del modelo durante Accelerator.save_state se dividirán en n archivos para cada sub-división del modelo. Para volver a fusionarlos en un único diccionario y cargarlos en el modelo más tarde, después del entrenamiento, puedes usar la utilidad merge_weights:

from accelerate.utils import merge_fsdp_weights

# Our weights are saved usually in a `pytorch_model_fsdp_{model_number}` folder
merge_fsdp_weights("pytorch_model_fsdp_0", "output_path", safe_serialization=True)

La salida final se guardará en model.safetensors o en pytorch_model.bin (si se pasa safe_serialization=False).

También se puede llamar usando la CLI:

accelerate merge-weights pytorch_model_fsdp_0/ output_path

Correspondencia entre las estrategias de particionamiento de FSDP y las etapas ZeRO de DeepSpeed

• FULL_SHARD corresponde a ZeRO Stage-3 de DeepSpeed. Fragmenta los estados del optimizador, los gradientes y los parámetros.
• SHARD_GRAD_OP corresponde a ZeRO Stage-2 de DeepSpeed. Fragmenta los estados del optimizador y los gradientes.
• NO_SHARD corresponde a ZeRO Stage-0. No hay particionamiento, cada GPU tiene una copia completa del modelo, los estados del optimizador y los gradientes.
• HYBRID_SHARD corresponde a ZeRO++ Stage-3, donde zero_hpz_partition_size=<num_gpus_per_node>. En este caso, se fragmentan los estados del optimizador, los gradientes y los parámetros dentro de cada nodo, mientras que cada nodo conserva una copia completa.

Algunas consideraciones a tener en cuenta

• En caso de varios modelos, pase los optimizadores a la llamada prepare en el mismo orden que los modelos correspondientes; de lo contrario, accelerator.save_state() y accelerator.load_state() provocarán un comportamiento incorrecto o inesperado.
• Esta función no es compatible con --predict_with_generate en el script run_translation.py de la biblioteca Transformers.

Para tener más control, los usuarios pueden aprovechar el FullyShardedDataParallelPlugin. Después de crear una instancia de esta clase, los usuarios pueden pasarla a la instancia de la clase Accelerator. Para obtener más información sobre estas opciones, consulte el código de FullyShardedDataParallel de PyTorch.

Para quienes estén interesados en las similitudes y diferencias entre FSDP y DeepSpeed, ¡consulten la guía de conceptos aquí!