Acelera el entrenamiento de modelos grandes utilizando DeepSpeed

Entrena modelos grandes más rápido con DeepSpeed.

En esta publicación veremos cómo podemos aprovechar la biblioteca Accelerate para entrenar modelos grandes, lo cual permite a los usuarios aprovechar las características de ZeRO de DeeSpeed.

¿Cansado de los errores de falta de memoria (OOM) al intentar entrenar modelos grandes? Tenemos la solución para ti. Los modelos grandes son muy eficientes [1], pero difíciles de entrenar con el hardware disponible. Para aprovechar al máximo el hardware disponible para entrenar modelos grandes, se puede utilizar el Paralelismo de Datos utilizando ZeRO – Optimizador de Redundancia Cero [2].

A continuación se muestra una breve descripción del Paralelismo de Datos utilizando ZeRO con un diagrama de esta publicación del blog

(Fuente: enlace)

a. Etapa 1: Distribuye los estados del optimizador entre los trabajadores/GPUs paralelos de datos

b. Etapa 2: Distribuye los estados del optimizador + gradientes entre los trabajadores/GPUs paralelos de datos

c. Etapa 3: Distribuye los estados del optimizador + gradientes + parámetros del modelo entre los trabajadores/GPUs paralelos de datos

d. Descarga del Optimizador: Descarga los gradientes + estados del optimizador a la CPU/Disco basándose en ZeRO Etapa 2

e. Descarga de Parámetros: Descarga los parámetros del modelo a la CPU/Disco basándose en ZeRO Etapa 3

En esta publicación del blog veremos cómo aprovechar el Paralelismo de Datos utilizando ZeRO con Accelerate. DeepSpeed, FairScale y PyTorch FullyShardedDataParallel (FSDP) han implementado las ideas principales del artículo de ZeRO. Estos ya se han integrado en 🤗 transformers Trainer y 🤗 accelerate, acompañados de excelentes blogs como Fit More and Train Faster With ZeRO via DeepSpeed and FairScale [4] y Accelerate Large Model Training using PyTorch Fully Sharded Data Parallel [5]. Dejamos la explicación de lo que ocurre detrás de escena en esos blogs y nos enfocamos principalmente en aprovechar DeepSpeed ZeRO utilizando Accelerate.

Configuración de Hardware: 2X24GB NVIDIA Titan RTX GPUs. 60GB RAM.

Veremos la tarea de ajustar un modelo solo codificador para clasificación de texto. Utilizaremos el modelo preentrenado microsoft/deberta-v2-xlarge-mnli (900M parámetros) para ajustarlo en el conjunto de datos MRPC GLUE.

El código está disponible aquí en run_cls_no_trainer.py. Es similar al ejemplo oficial de clasificación de texto aquí, con la adición de la lógica para medir el tiempo de entrenamiento y evaluación. Comparemos el rendimiento entre el Paralelismo de Datos Distribuido (DDP) y DeepSpeed ZeRO Etapa-2 en una configuración Multi-GPU.

Para habilitar DeepSpeed ZeRO Etapa-2 sin realizar ningún cambio en el código, ejecuta accelerate config y aprovecha el complemento Accelerate DeepSpeed.

Ejemplo de Complemento DeepSpeed ZeRO Etapa-2

compute_environment: LOCAL_MACHINE
deepspeed_config:
 gradient_accumulation_steps: 1
 gradient_clipping: 1.0
 offload_optimizer_device: none
 offload_param_device: none
 zero3_init_flag: false
 zero_stage: 2
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false

Ahora, ejecuta el siguiente comando para el entrenamiento:

accelerate launch run_cls_no_trainer.py \
  --model_name_or_path "microsoft/deberta-v2-xlarge-mnli" \
  --task_name "mrpc" \
  --ignore_mismatched_sizes \
  --max_length 128 \
  --per_device_train_batch_size 40 \
  --learning_rate 2e-5 \
  --num_train_epochs 3 \
  --output_dir "/tmp/mrpc/deepspeed_stage2/" \
  --with_tracking \
  --report_to "wandb" \

En nuestra configuración de un solo nodo con múltiples GPU, el tamaño de lote máximo que admite DDP sin errores de falta de memoria es 8. En cambio, DeepSpeed ZeRO Etapa 2 permite un tamaño de lote de 40 sin encontrar errores de falta de memoria. Por lo tanto, DeepSpeed permite ajustar 5 veces más datos por GPU en comparación con DDP. A continuación se muestra la captura de pantalla de los gráficos de la ejecución en wandb junto con una tabla de referencia comparando DDP vs DeepSpeed.


Tabla 1: Pruebas de rendimiento de DeepSpeed ZeRO Stage-2 en el modelo DeBERTa-XL (900M)


Con este tamaño de lote más grande, observamos una aceleración de ~ 3.5X en el tiempo total de entrenamiento sin ninguna disminución en las métricas de rendimiento, todo esto sin cambiar ningún código. ¡Yeah! 🤗.

Para poder ajustar más opciones, deberás usar un archivo de configuración de DeepSpeed y hacer cambios mínimos en el código. Veamos cómo hacer esto.

Primero, veremos la tarea de ajustar finamente un modelo de secuencia a secuencia para entrenar nuestro propio Chatbot. Específicamente, ajustaremos finamente facebook/blenderbot-400M-distill en el conjunto de datos smangrul/MuDoConv (Conversación de múltiples dominios). El conjunto de datos contiene conversaciones de 10 fuentes de datos diferentes que cubren personas, fundamentos en contextos emocionales específicos, orientados a objetivos (por ejemplo, reserva de restaurantes) y temas generales de Wikipedia (por ejemplo, Cricket).

El código está disponible aquí run_seq2seq_no_trainer.py. La práctica actual para medir de manera efectiva la Participación y la Humanidad de los Chatbots es mediante evaluaciones humanas que son costosas [6]. Como tal, para este ejemplo, la métrica que se está rastreando es la puntuación BLEU (que no es ideal pero es la métrica convencional para este tipo de tareas). Se puede adaptar el código para entrenar modelos T5 más grandes si se tiene acceso a GPUs que admitan la precisión bfloat16 de lo contrario se obtendrán valores de pérdida NaN. Realizaremos una prueba rápida en 10000 muestras de entrenamiento y 1000 muestras de evaluación, ya que estamos interesados en DeepSpeed vs DDP.

Aprovecharemos la configuración de DeepSpeed Zero Stage-2 zero2_config_accelerate.json (que se muestra a continuación) para el entrenamiento. Para obtener información detallada sobre las diferentes características de configuración, consulta la documentación de DeeSpeed.

{
    "fp16": {
        "enabled": "true",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 15,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "weight_decay": "auto",
            "torch_adam": true,
            "adam_w_mode": true
        }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto",
            "total_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 2,
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

Para habilitar DeepSpeed ZeRO Stage-2 con la configuración anterior, ejecuta accelerate config y proporciona la ruta del archivo de configuración cuando se te pregunte. Para obtener más detalles, consulta la documentación oficial de 🤗 accelerate para el archivo de configuración de DeepSpeed.

Ejemplo de archivo de configuración de DeepSpeed ZeRO Stage-2

compute_environment: LOCAL_MACHINE
deepspeed_config:
 deepspeed_config_file: /ruta/a/zero2_config_accelerate.json
 zero3_init_flag: false
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false

Ahora, ejecuta el siguiente comando para el entrenamiento:

accelerate launch run_seq2seq_no_trainer.py \
    --dataset_name "smangrul/MuDoConv" \
    --max_source_length 128 \
    --source_prefix "chatbot: " \
    --max_target_length 64 \
    --val_max_target_length 64 \
    --val_min_target_length 20 \
    --n_val_batch_generations 5 \
    --n_train 10000 \
    --n_val 1000 \
    --pad_to_max_length \
    --num_beams 10 \
    --model_name_or_path "facebook/blenderbot-400M-distill" \
    --per_device_train_batch_size 200 \
    --per_device_eval_batch_size 100 \
    --learning_rate 1e-6 \
    --weight_decay 0.0 \
    --num_train_epochs 1 \
    --gradient_accumulation_steps 1 \
    --num_warmup_steps 100 \
    --output_dir "/tmp/deepspeed_zero_stage2_accelerate_test" \
    --seed 25 \
    --logging_steps 100 \
    --with_tracking \
    --report_to "wandb" \
    --report_name "blenderbot_400M_finetuning"

Cuando se utiliza la configuración de DeepSpeed, si el usuario ha especificado optimizer y scheduler en la configuración, el usuario deberá utilizar accelerate.utils.DummyOptim y accelerate.utils.DummyScheduler. Esos son los únicos cambios menores que el usuario debe hacer. A continuación mostramos un ejemplo de los cambios mínimos requeridos al utilizar la configuración de DeepSpeed:

- optimizer = torch.optim.Adam(optimizer_grouped_parameters, lr=args.learning_rate)
+ optimizer = accelerate.utils.DummyOptim(optimizer_grouped_parameters, lr=args.learning_rate)

- lr_scheduler = get_scheduler(
-     name=args.lr_scheduler_type,
-     optimizer=optimizer,
-     num_warmup_steps=args.num_warmup_steps,
-     num_training_steps=args.max_train_steps,
- )

+ lr_scheduler = accelerate.utils.DummyScheduler(
+     optimizer, total_num_steps=args.max_train_steps, warmup_num_steps=args.num_warmup_steps
+ )

Tabla 2: Evaluación del rendimiento de DeepSpeed ZeRO Stage-2 en el modelo BlenderBot (400M)

En nuestra configuración de un solo nodo con múltiples GPU, el tamaño máximo de lote que DDP admite sin errores de falta de memoria (OOM) es 100. En cambio, DeepSpeed Zero-Stage 2 permite un tamaño de lote de 200 sin errores de falta de memoria. Por lo tanto, DeepSpeed permite ajustar 2X más datos por GPU en comparación con DDP. Observamos un aumento de velocidad de ~ 1.44X en el entrenamiento y ~ 1.23X en la evaluación, ya que podemos ajustar más datos en el mismo hardware disponible. Como este modelo tiene el tamaño de VoAGI, el aumento de velocidad no es tan emocionante, pero esto mejorará con modelos más grandes. Puedes chatear con el Chatbot entrenado usando todos los datos en 🤗 Space smangrul/Chat-E. Puedes darle una personalidad al bot, una conversación de referencia a una emoción particular, usarlo en tareas orientadas a objetivos o de manera libre. A continuación se muestra una conversación divertida con el chatbot 💬. Puedes encontrar capturas de pantalla de más conversaciones utilizando diferentes contextos aquí.


Desvío de CPU/Disco para habilitar el entrenamiento de modelos enormes que no caben en la memoria de la GPU

En una sola GPU NVIDIA Titan RTX de 24 GB, no se puede entrenar el modelo GPT-XL (1,5B de parámetros) incluso con un tamaño de lote de 1. Veremos cómo podemos utilizar DeepSpeed ZeRO Stage-3 con el desvío de CPU de los estados del optimizador, gradientes y parámetros para entrenar el modelo GPT-XL.

Aprovecharemos la configuración de desvío de CPU Zero Stage-3 de DeepSpeed zero3_offload_config_accelerate.json (que se muestra a continuación) para el entrenamiento. El resto del proceso de uso de la configuración con 🤗 accelerate es similar al experimento anterior.

{
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "weight_decay": "auto"
        }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto",
            "total_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "sub_group_size": 1e9,
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

Ejemplo de archivo de configuración de DeepSpeed ZeRO Stage-3 CPU Offload

compute_environment: LOCAL_MACHINE
deepspeed_config:
 deepspeed_config_file: /ruta/a/zero3_offload_config_accelerate.json
 zero3_init_flag: true
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false

Ahora, ejecuta el siguiente comando para el entrenamiento:

accelerate launch run_clm_no_trainer.py \
--config_name "gpt2-xl" \
--tokenizer_name "gpt2-xl" \
--dataset_name "wikitext" \
--dataset_config_name "wikitext-2-raw-v1" \
--block_size 128 \
--output_dir "/tmp/clm_deepspeed_stage3_offload__accelerate" \
--learning_rate 5e-4 \
--per_device_train_batch_size 16 \
--per_device_eval_batch_size 1 \
--num_train_epochs 1 \
--with_tracking \
--report_to "wandb"\

Tabla 3: Evaluación de DeepSpeed ZeRO Stage-3 CPU Offload en el modelo GPT-XL (1.5B)


DDP resultará en un error de OOM incluso con un tamaño de lote de 1. Por otro lado, con DeepSpeed ZeRO Stage-3 CPU offload, podemos entrenar con un tamaño de lote de 16.

Finalmente, por favor, recuerda que, 🤗 Accelerate solo integra DeepSpeed, por lo tanto, si tienes algún problema o pregunta con respecto al uso de DeepSpeed, por favor, presenta un problema en el repositorio de GitHub de DeepSpeed.

[1] Entrena grande, luego comprime: Repensando el tamaño del modelo para un entrenamiento y una inferencia eficientes de los Transformers

[2] ZeRO: Optimizaciones de memoria hacia el entrenamiento de modelos de billones de parámetros

[3] DeepSpeed: Entrenamiento de modelos a escala extrema para todos – Microsoft Research

[4] Ajusta más y entrena más rápido con ZeRO a través de DeepSpeed y FairScale

[5] Acelera el entrenamiento de modelos grandes utilizando PyTorch Fully Sharded Data Parallel

[6] Recetas para construir un chatbot de dominio abierto

We will continue to update Zepes; if you have any questions or suggestions, please contact us!

Share:

Was this article helpful?

93 out of 132 found this helpful

Discover more

Inteligencia Artificial

Investigadores de Apple proponen un nuevo modelo de descomposición de tensores para el filtrado colaborativo con retroalimentación implícita

La capacidad para inferir las preferencias del usuario a partir de comportamientos pasados es crucial para ofrecer su...

Inteligencia Artificial

Descifrando el comportamiento colectivo Cómo la inferencia bayesiana activa impulsa los movimientos naturales de los grupos de animales

El fenómeno del movimiento colectivo en animales observado en actividades como los enjambres de langostas, los cardúm...

Inteligencia Artificial

Los modelos base Llama 2 de Meta ahora están disponibles en Amazon SageMaker JumpStart

Hoy, nos complace anunciar que los modelos base Llama 2 desarrollados por Meta están disponibles para los clientes a ...

Noticias de Inteligencia Artificial

Utilice AWS CDK para implementar configuraciones del ciclo de vida de Amazon SageMaker Studio.

Amazon SageMaker Studio es el primer entorno de desarrollo integrado (IDE) completamente para machine learning (ML). ...

Noticias de Inteligencia Artificial

Los Nano-Tatuajes No Necesitan Baterías ni Cables

Los sensores de nano-tatuajes basados en la retrodispersión pueden comunicarse con dispositivos cercanos sin necesida...