DeepSpeed

DeepSpeed es una biblioteca diseñada para la velocidad y la escala en el entrenamiento distribuido de modelos grandes con miles de millones de parámetros. En su núcleo se encuentra el Zero Redundancy Optimizer (ZeRO), que fragmenta los estados del optimizador (ZeRO-1), los gradientes (ZeRO-2) y los parámetros (ZeRO-3) entre procesos de paralelismo de datos. Esto reduce drásticamente el uso de memoria, lo que permite escalar el entrenamiento a modelos con miles de millones de parámetros. Para desbloquear aún más eficiencia de memoria, ZeRO-Offload reduce el cómputo y la memoria de la GPU aprovechando los recursos de la CPU durante la optimización.

Ambas funciones son compatibles con 🤗 Accelerate y puedes usarlas con 🤗 PEFT.

Compatibilidad con cuantización bitsandbytes + LoRA

A continuación se muestra una tabla que resume la compatibilidad entre el LoRA de PEFT, la biblioteca bitsandbytes y las etapas Zero de DeepSpeed en relación con el ajuste fino. DeepSpeed Zero-1 y 2 no tendrán efecto en la inferencia, ya que la etapa 1 fragmenta los estados del optimizador y la etapa 2 fragmenta los estados del optimizador y los gradientes:

Etapa de DeepSpeed	¿Es compatible?
Zero-1	🟢
Zero-2	🟢
Zero-3	🟢

Para DeepSpeed Stage 3 + QLoRA, consulte la sección Use PEFT QLoRA and DeepSpeed with ZeRO3 for finetuning large models on multiple GPUs a continuación.

Para confirmar estas observaciones, ejecutamos los scripts de ejemplo oficiales de SFT (Supervised Fine-tuning) de la librería Transformers Reinforcement Learning (TRL) utilizando QLoRA + PEFT y las configuraciones de accelerate disponibles aquí. Ejecutamos estos experimentos en una GPU 2x NVIDIA T4.

Usa PEFT y DeepSpeed con ZeRO3 para el ajuste fino de modelos grandes en varios dispositivos y varios nodos

Esta sección de la guía te ayudará a aprender cómo usar nuestro script de entrenamiento de DeepSpeed para realizar SFT. Configurarás el script para hacer SFT (ajuste fino supervisado) del modelo Llama-70B con LoRA y ZeRO-3 en 8 GPUs H100 de 80 GB en una sola máquina. Puedes configurarlo para escalar a múltiples máquinas cambiando la configuración de accelerate.

Configuración

Comience ejecutando el siguiente comando para crear un archivo de configuración de DeepSpeed con 🤗 Accelerate. La bandera --config_file le permite guardar el archivo de configuración en una ubicación específica; de lo contrario, se guarda como un archivo default_config.yaml en la caché de 🤗 Accelerate.

El archivo de configuración se utiliza para establecer las opciones predeterminadas al iniciar el script de entrenamiento.

accelerate config --config_file deepspeed_config.yaml

Se te harán algunas preguntas sobre tu configuración y deberás configurar los siguientes argumentos. En este ejemplo, usarás ZeRO-3, así que asegúrate de seleccionar esas opciones.

`zero_stage`: [0] Disabled, [1] optimizer state partitioning, [2] optimizer+gradient state partitioning and [3] optimizer+gradient+parameter partitioning
`gradient_accumulation_steps`: Number of training steps to accumulate gradients before averaging and applying them. Pass the same value as you would pass via cmd argument else you will encounter mismatch error.
`gradient_clipping`: Enable gradient clipping with value. Don't set this as you will be passing it via cmd arguments.
`offload_optimizer_device`: [none] Disable optimizer offloading, [cpu] offload optimizer to CPU, [nvme] offload optimizer to NVMe SSD. Only applicable with ZeRO >= Stage-2. Set this as `none` as don't want to enable offloading.
`offload_param_device`: [none] Disable parameter offloading, [cpu] offload parameters to CPU, [nvme] offload parameters to NVMe SSD. Only applicable with ZeRO Stage-3. Set this as `none` as don't want to enable offloading.
`zero3_init_flag`: Decides whether to enable `deepspeed.zero.Init` for constructing massive models. Only applicable with ZeRO Stage-3. Set this to `True`.
`zero3_save_16bit_model`: Decides whether to save 16-bit model weights when using ZeRO Stage-3. Set this to `True`.
`mixed_precision`: `no` for FP32 training, `fp16` for FP16 mixed-precision training and `bf16` for BF16 mixed-precision training. Set this to `True`.

Una vez hecho esto, la configuración correspondiente debería verse como la siguiente y puedes encontrarla en la carpeta config en deepspeed_config.yaml:

compute_environment: LOCAL_MACHINE                                                                                                                                           
debug: false
deepspeed_config:
  deepspeed_multinode_launcher: standard
  gradient_accumulation_steps: 4
  offload_optimizer_device: none
  offload_param_device: none
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
distributed_type: DEEPSPEED
downcast_bf16: 'no'
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

Comando de lanzamiento

El comando de lanzamiento está disponible en run_peft_deepspeed.sh y también se muestra a continuación:

accelerate launch --config_file "configs/deepspeed_config.yaml"  train.py \
--seed 100 \
--model_name_or_path "meta-llama/Llama-2-70b-hf" \
--dataset_name "smangrul/ultrachat-10k-chatml" \
--chat_template_format "chatml" \
--add_special_tokens False \
--append_concat_token False \
--splits "train,test" \
--max_seq_len 2048 \
--num_train_epochs 1 \
--logging_steps 5 \
--log_level "info" \
--logging_strategy "steps" \
--eval_strategy "epoch" \
--save_strategy "epoch" \
--push_to_hub \
--hub_private_repo True \
--hub_strategy "every_save" \
--bf16 True \
--packing True \
--learning_rate 1e-4 \
--lr_scheduler_type "cosine" \
--weight_decay 1e-4 \
--warmup_steps 0 \
--max_grad_norm 1.0 \
--output_dir "llama-sft-lora-deepspeed" \
--per_device_train_batch_size 8 \
--per_device_eval_batch_size 8 \
--gradient_accumulation_steps 4 \
--gradient_checkpointing True \
--use_reentrant False \
--dataset_text_field "content" \
--use_flash_attn True \
--use_peft_lora True \
--lora_r 8 \
--lora_alpha 16 \
--lora_dropout 0.1 \
--lora_target_modules "all-linear" \
--use_4bit_quantization False

Tenga en cuenta que estamos usando LoRA con rank=8, alpha=16 y apuntando a todas las capas lineales. Estamos pasando el archivo de configuración de DeepSpeed y ajustando el modelo Llama de 70B en un subconjunto del conjunto de datos ultrachat.

Las partes importantes

Vamos a profundizar un poco más en el script para que puedas ver qué está pasando y entender cómo funciona.

Lo primero que hay que saber es que el script utiliza DeepSpeed para el entrenamiento distribuido, ya que se ha pasado la configuración de DeepSpeed. La clase SFTTrainer se encarga de todo el trabajo pesado de crear el modelo PEFT utilizando la configuración de peft que se pasa. Después, cuando se llama a trainer.train(), SFTTrainer utiliza internamente 🤗 Accelerate para preparar el modelo, el optimizador y el entrenador, usando la configuración de DeepSpeed para crear el motor de DeepSpeed, que luego se entrena. El fragmento de código principal es el siguiente:

# trainer
trainer = SFTTrainer(
    model=model,
    processing_class=tokenizer,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    peft_config=peft_config,
)
trainer.accelerator.print(f"{trainer.model}")

# train
checkpoint = None
if training_args.resume_from_checkpoint is not None:
    checkpoint = training_args.resume_from_checkpoint
trainer.train(resume_from_checkpoint=checkpoint)

# saving final model
trainer.save_model()

Uso de memoria

En el ejemplo anterior, la memoria consumida por GPU es de 64 GB (80 %), como se ve en la captura de pantalla siguiente:

GPU memory usage for the training run

Más recursos

También puede consultar esta publicación de blog Falcon 180B Finetuning using 🤗 PEFT and DeepSpeed sobre cómo ajustar el modelo Falcon de 180B en 16 GPUs A100 en 2 máquinas.

Usa PEFT QLoRA y DeepSpeed con ZeRO3 para el ajuste fino de modelos grandes en varias GPU

En esta sección, veremos cómo usar QLoRA y DeepSpeed Stage-3 para el ajuste fino del modelo llama de 70B en 2X40GB GPUs. Para ello, primero necesitamos bitsandbytes>=0.43.3, accelerate>=1.0.1, transformers>4.44.2, trl>0.11.4 y peft>0.13.0. Debemos establecer zero3_init_flag en true al usar la configuración de Accelerate. A continuación se muestra la configuración que se puede encontrar en deepspeed_config_z3_qlora.yaml:

compute_environment: LOCAL_MACHINE                                                                                                                                           
debug: false
deepspeed_config:
  deepspeed_multinode_launcher: standard
  offload_optimizer_device: none
  offload_param_device: none
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
distributed_type: DEEPSPEED
downcast_bf16: 'no'
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

El comando de lanzamiento se muestra a continuación y está disponible en run_peft_qlora_deepspeed_stage3.sh:

accelerate launch --config_file "configs/deepspeed_config_z3_qlora.yaml"  train.py \
--seed 100 \
--model_name_or_path "meta-llama/Llama-2-70b-hf" \
--dataset_name "smangrul/ultrachat-10k-chatml" \
--chat_template_format "chatml" \
--add_special_tokens False \
--append_concat_token False \
--splits "train,test" \
--max_seq_len 2048 \
--num_train_epochs 1 \
--logging_steps 5 \
--log_level "info" \
--logging_strategy "steps" \
--eval_strategy "epoch" \
--save_strategy "epoch" \
--push_to_hub \
--hub_private_repo True \
--hub_strategy "every_save" \
--bf16 True \
--packing True \
--learning_rate 1e-4 \
--lr_scheduler_type "cosine" \
--weight_decay 1e-4 \
--warmup_steps 0 \
--max_grad_norm 1.0 \
--output_dir "llama-sft-qlora-dsz3" \
--per_device_train_batch_size 2 \
--per_device_eval_batch_size 2 \
--gradient_accumulation_steps 2 \
--gradient_checkpointing True \
--use_reentrant True \
--dataset_text_field "content" \
--use_flash_attn True \
--use_peft_lora True \
--lora_r 8 \
--lora_alpha 16 \
--lora_dropout 0.1 \
--lora_target_modules "all-linear" \
--use_4bit_quantization True \
--use_nested_quant True \
--bnb_4bit_compute_dtype "bfloat16" \
--bnb_4bit_quant_storage_dtype "bfloat16"

Tenga en cuenta el nuevo argumento que se pasa, bnb_4bit_quant_storage_dtype, que denota el tipo de datos para el empaquetado de los parámetros de 4 bits. Por ejemplo, cuando se establece en bfloat16, 32/4 = 8 parámetros de 4 bits se empaquetan juntos tras la cuantización.

En cuanto al código de entrenamiento, los cambios importantes en el código son:

...

bnb_config = BitsAndBytesConfig(
    load_in_4bit=args.use_4bit_quantization,
    bnb_4bit_quant_type=args.bnb_4bit_quant_type,
    bnb_4bit_compute_dtype=compute_dtype,
    bnb_4bit_use_double_quant=args.use_nested_quant,
+   bnb_4bit_quant_storage=quant_storage_dtype,
)

...

model = AutoModelForCausalLM.from_pretrained(
    args.model_name_or_path,
    quantization_config=bnb_config,
    trust_remote_code=True,
    attn_implementation="flash_attention_2" if args.use_flash_attn else "eager",
+   dtype=quant_storage_dtype or torch.float32,
)

Tenga en cuenta que el dtype de AutoModelForCausalLM es el mismo que el tipo de datos bnb_4bit_quant_storage. Eso es todo. Todo lo demás lo manejan Trainer y TRL.

Uso de memoria

En el ejemplo anterior, la memoria consumida por GPU es de 36.6 GB. Por lo tanto, lo que requería 8 GPUs de 80GB con DeepSpeed Stage 3+LoRA y un par de GPUs de 80GB con DDP+QLoRA ahora requiere 2 GPUs de 40GB. Esto hace que el ajuste fino de modelos grandes sea más accesible.

Usa PEFT y DeepSpeed con ZeRO3 y CPU Offloading para el ajuste fino de modelos grandes en una sola GPU

Esta sección de la guía te ayudará a aprender cómo usar nuestro script de entrenamiento de DeepSpeed. Configurarás el script para entrenar un modelo grande para generación condicional con ZeRO-3 y CPU Offload.

💡 Para ayudarte a empezar, consulta nuestros scripts de ejemplo de entrenamiento para modelado de lenguaje causal y generación condicional. Puedes adaptar estos scripts para tus propias aplicaciones o incluso usarlos tal cual si tu tarea es similar a la de los scripts.

Configuración

El archivo de configuración se utiliza para establecer las opciones predeterminadas al iniciar el script de entrenamiento.

accelerate config --config_file ds_zero3_cpu.yaml

Se te harán algunas preguntas sobre tu configuración y deberás configurar los siguientes argumentos. En este ejemplo, usarás ZeRO-3 junto con CPU-Offload, así que asegúrate de seleccionar esas opciones.

`zero_stage`: [0] Disabled, [1] optimizer state partitioning, [2] optimizer+gradient state partitioning and [3] optimizer+gradient+parameter partitioning
`gradient_accumulation_steps`: Number of training steps to accumulate gradients before averaging and applying them.
`gradient_clipping`: Enable gradient clipping with value.
`offload_optimizer_device`: [none] Disable optimizer offloading, [cpu] offload optimizer to CPU, [nvme] offload optimizer to NVMe SSD. Only applicable with ZeRO >= Stage-2.
`offload_param_device`: [none] Disable parameter offloading, [cpu] offload parameters to CPU, [nvme] offload parameters to NVMe SSD. Only applicable with ZeRO Stage-3.
`zero3_init_flag`: Decides whether to enable `deepspeed.zero.Init` for constructing massive models. Only applicable with ZeRO Stage-3.
`zero3_save_16bit_model`: Decides whether to save 16-bit model weights when using ZeRO Stage-3.
`mixed_precision`: `no` for FP32 training, `fp16` for FP16 mixed-precision training and `bf16` for BF16 mixed-precision training.

Un ejemplo de archivo de configuración podría ser el siguiente. Lo más importante a tener en cuenta es que zero_stage está establecido en 3, y offload_optimizer_device y offload_param_device están establecidos en cpu.

compute_environment: LOCAL_MACHINE
deepspeed_config:
  gradient_accumulation_steps: 1
  gradient_clipping: 1.0
  offload_optimizer_device: cpu
  offload_param_device: cpu
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
distributed_type: DEEPSPEED
downcast_bf16: 'no'
dynamo_backend: 'NO'
fsdp_config: {}
machine_rank: 0
main_training_function: main
megatron_lm_config: {}
mixed_precision: 'no'
num_machines: 1
num_processes: 1
rdzv_backend: static
same_network: true
use_cpu: false

Las partes importantes

Vamos a profundizar un poco más en el script para que puedas ver qué está pasando y entender cómo funciona.

Dentro de la función main, el script crea una clase Accelerator para inicializar todos los requisitos necesarios para el entrenamiento distribuido.

💡 Puedes cambiar el modelo y el conjunto de datos dentro de la función main. Si el formato de tu conjunto de datos es diferente al del script, es posible que también necesites escribir tu propia función de preprocesamiento.

El script también crea una configuración para el método de 🤗 PEFT que estás utilizando, que en este caso es LoRA. LoraConfig especifica el tipo de tarea y parámetros importantes como la dimensión de las matrices de bajo rango, el factor de escalado de las matrices y la probabilidad de dropout de las capas de LoRA. Si deseas utilizar un método de 🤗 PEFT diferente, asegúrate de reemplazar LoraConfig con la clase correspondiente.

 def main():
+    accelerator = Accelerator()
     model_name_or_path = "facebook/bart-large"
     dataset_name = "twitter_complaints"
+    peft_config = LoraConfig(
         task_type=TaskType.SEQ_2_SEQ_LM, inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1
     )

A lo largo del script, verás las funciones main_process_first y wait_for_everyone, que ayudan a controlar y sincronizar cuándo se ejecutan los procesos.

La función get_peft_model() toma un modelo base y el peft_config que preparaste anteriormente para crear un PeftModel:

  model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)
+ model = get_peft_model(model, peft_config)

Pasa todos los objetos de entrenamiento relevantes a prepare de 🤗 Accelerate, que se asegura de que todo esté listo para el entrenamiento:

model, train_dataloader, eval_dataloader, test_dataloader, optimizer, lr_scheduler = accelerator.prepare(
    model, train_dataloader, eval_dataloader, test_dataloader, optimizer, lr_scheduler
)

El siguiente fragmento de código comprueba si el plugin de DeepSpeed se utiliza en el Accelerator, y si el plugin existe, entonces se verifica si se está usando ZeRO-3. Esta bandera condicional se utiliza al llamar a la función generate durante la inferencia para sincronizar las GPUs cuando los parámetros del modelo están fragmentados:

is_ds_zero_3 = False
if getattr(accelerator.state, "deepspeed_plugin", None):
    is_ds_zero_3 = accelerator.state.deepspeed_plugin.zero_stage == 3

Dentro del bucle de entrenamiento, el habitual loss.backward() se reemplaza por backward de 🤗 Accelerate, que utiliza el método backward() correcto según tu configuración:

  for epoch in range(num_epochs):
      with TorchTracemalloc() as tracemalloc:
          model.train()
          total_loss = 0
          for step, batch in enumerate(tqdm(train_dataloader)):
              outputs = model(**batch)
              loss = outputs.loss
              total_loss += loss.detach().float()
+             accelerator.backward(loss)
              optimizer.step()
              lr_scheduler.step()
              optimizer.zero_grad()

¡Eso es todo! El resto del script gestiona el bucle de entrenamiento, la evaluación e incluso lo sube al Hub por ti.

Entrenar

Ejecuta el siguiente comando para iniciar el script de entrenamiento. Antes, guardaste el archivo de configuración en ds_zero3_cpu.yaml, por lo que necesitas pasar la ruta al lanzador con el argumento --config_file de la siguiente manera:

accelerate launch --config_file ds_zero3_cpu.yaml examples/conditional_generation/peft_lora_seq2seq_accelerate_ds_zero3_offload.py

Verás algunos registros de salida que rastrean el uso de memoria durante el entrenamiento y, una vez que este se completa, el script devuelve la precisión y compara las predicciones con las etiquetas:

GPU Memory before entering the train : 1916
GPU Memory consumed at the end of the train (end-begin): 66
GPU Peak Memory consumed during the train (max-begin): 7488
GPU Total Peak Memory consumed during the train (max): 9404
CPU Memory before entering the train : 19411
CPU Memory consumed at the end of the train (end-begin): 0
CPU Peak Memory consumed during the train (max-begin): 0
CPU Total Peak Memory consumed during the train (max): 19411
epoch=4: train_ppl=tensor(1.0705, device='cuda:0') train_epoch_loss=tensor(0.0681, device='cuda:0')
100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:27<00:00,  3.92s/it]
GPU Memory before entering the eval : 1982
GPU Memory consumed at the end of the eval (end-begin): -66
GPU Peak Memory consumed during the eval (max-begin): 672
GPU Total Peak Memory consumed during the eval (max): 2654
CPU Memory before entering the eval : 19411
CPU Memory consumed at the end of the eval (end-begin): 0
CPU Peak Memory consumed during the eval (max-begin): 0
CPU Total Peak Memory consumed during the eval (max): 19411
accuracy=100.0
eval_preds[:10]=['no complaint', 'no complaint', 'complaint', 'complaint', 'no complaint', 'no complaint', 'no complaint', 'complaint', 'complaint', 'no complaint']
dataset['train'][label_column][:10]=['no complaint', 'no complaint', 'complaint', 'complaint', 'no complaint', 'no complaint', 'no complaint', 'complaint', 'complaint', 'no complaint']

Precauciones

La fusión al usar PEFT y DeepSpeed no se admite actualmente y generará un error.
Al usar el descargado a CPU, las principales ganancias de usar PEFT para reducir los estados del optimizador y los gradientes al tamaño de los pesos del adaptador se materializarían en la RAM de la CPU y no habría ahorros con respecto a la memoria de la GPU.
DeepSpeed Stage 3 y qlora, cuando se usan con CPU offloading, provocan un mayor uso de memoria GPU en comparación con deshabilitar el CPU offloading.

💡 Cuando tienes código que requiere la fusión (y desfusión) de pesos, intenta recopilar manualmente los parámetros con DeepSpeed Zero-3 de antemano:
import deepspeed

is_ds_zero_3 = ... # check if Zero-3

with deepspeed.zero.GatheredParameters(list(model.parameters()), enabled= is_ds_zero_3):
    model.merge_adapter()
    # do whatever is needed, then unmerge in the same context if unmerging is required
    ...
    model.unmerge_adapter()

Update on GitHub