Paralelismo de datos totalmente fragmentado

Fully sharded data parallel (FSDP) se desarrolló para el entrenamiento distribuido de grandes modelos preentrenados de hasta 1T de parámetros. FSDP lo logra fragmentando los parámetros del modelo, los gradientes y los estados del optimizador entre los procesos de paralelismo de datos, y también puede descargar los parámetros fragmentados del modelo a la CPU. La eficiencia de memoria que ofrece FSDP permite escalar el entrenamiento a tamaños de lote o de modelo más grandes.

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

Usar PEFT y FSDP

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 FSDP 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 FSDP 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 fsdp_config.yaml

Se te harán algunas preguntas sobre tu configuración y deberás configurar los siguientes argumentos. En este ejemplo, responderás al cuestionario como se muestra en la imagen siguiente.

Creating Accelerate's config to use FSDP

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

compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch: BACKWARD_PRE
  fsdp_cpu_ram_efficient_loading: true
  fsdp_forward_prefetch: false
  fsdp_offload_params: false
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_use_orig_params: false
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_fsdp.sh y también se muestra a continuación:

accelerate launch --config_file "configs/fsdp_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-fsdp" \
--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 FSDP 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 FSDP para el entrenamiento distribuido, ya que se ha pasado la configuración de FSDP. 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(), Trainer utiliza internamente 🤗 Accelerate para preparar el modelo, el optimizador y el trainer, usando la configuración de FSDP para crear un modelo envuelto con FSDP 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}")
if model_args.use_peft_lora:
    # handle PEFT+FSDP case
    trainer.model.print_healthcheck()
    if getattr(trainer.accelerator.state, "fsdp_plugin", None):
        from peft.utils.other import fsdp_auto_wrap_policy

        fsdp_plugin = trainer.accelerator.state.fsdp_plugin
        fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(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
if trainer.is_fsdp_enabled:
    trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")
trainer.save_model()

Aquí, un punto principal a tener en cuenta actualmente al usar FSDP con PEFT es que use_orig_params debe ser False para lograr ahorros de memoria de GPU. Debido a use_orig_params=False, la política de auto wrap de FSDP debe cambiar para que los parámetros entrenables y no entrenables se envuelvan por separado. Esto se hace con el fragmento de código siguiente, que utiliza la función util fsdp_auto_wrap_policy de PEFT:

if getattr(trainer.accelerator.state, "fsdp_plugin", None):
    from peft.utils.other import fsdp_auto_wrap_policy

    fsdp_plugin = trainer.accelerator.state.fsdp_plugin
    fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(trainer.model)

Uso de memoria

En el ejemplo anterior, la memoria consumida por GPU es de 72-80 GB (90-98%), como se ve en la captura de pantalla siguiente. El ligero aumento de la memoria de GPU al final ocurre al guardar el modelo usando el tipo de state dict FULL_STATE_DICT en lugar de SHARDED_STATE_DICT, de modo que el modelo tenga pesos de adaptador que se puedan cargar normalmente con el método from_pretrained durante la inferencia:

GPU memory usage for the training run

Usa PEFT QLoRA y FSDP para el ajuste fino de modelos grandes en varias GPU

En esta sección, veremos cómo usar QLoRA y FSDP para el ajuste fino del modelo llama de 70B en 2X24GB de GPU. Answer.AI, en colaboración con bitsandbytes y Hugging Face 🤗, ha publicado el código de forma abierta que permite el uso de FSDP+QLoRA y explicó todo el proceso en su blogpost informativo You can now train a 70b language model at home. Esto ahora está integrado en el ecosistema de Hugging Face.

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 fsdp_cpu_ram_efficient_loading=true, fsdp_use_orig_params=false y fsdp_offload_params=true (descarga a CPU) al usar la configuración de Accelerate. Si no se utiliza el lanzador de accelerate, se puede establecer alternativamente la variable de entorno export FSDP_CPU_RAM_EFFICIENT_LOADING=true. Aquí utilizaremos la configuración de accelerate y a continuación se muestra la configuración que se puede encontrar en fsdp_config_qlora.yaml:

compute_environment: LOCAL_MACHINE                                                                                                                                           
debug: false                                                                                                                                                                 
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch: BACKWARD_PRE
  fsdp_cpu_ram_efficient_loading: true
  fsdp_forward_prefetch: false
  fsdp_offload_params: true
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_use_orig_params: false
machine_rank: 0
main_training_function: main
mixed_precision: 'no'
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_fsdp.sh:

accelerate launch --config_file "configs/fsdp_config_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-fsdp" \
--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"

Fíjense en el nuevo argumento que se pasa, bnb_4bit_quant_storage_dtype, que denota el tipo de datos para empaquetar los parámetros de 4 bits. Por ejemplo, cuando se establece en bfloat16, 16/4 = 4 parámetros de 4 bits se empaquetan juntos tras la cuantización. Al usar entrenamiento con precisión mixta con bfloat16, bnb_4bit_quant_storage_dtype puede ser bfloat16 para un ajuste fino puro en bfloat16, o float32 para precisión mixta automática (esto consume más memoria de GPU). Al usar entrenamiento con precisión mixta con float16, bnb_4bit_quant_storage_dtype debe establecerse en float32 para un entrenamiento con precisión mixta automática estable.

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 19.6 GB, mientras que el uso de RAM de CPU es de aproximadamente 107 GB. Al desactivar el offloading de CPU, el uso de memoria de GPU es de 35.6 GB/ GPU. Por lo tanto, lo que requería 16 GPUs de 80 GB para un finetuning completo, 8 GPUs de 80 GB con FSDP+LoRA y un par de GPUs de 80 GB con DDP+QLoRA, ahora requiere 2 GPUs de 24 GB. Esto hace que el finetuning de modelos grandes sea más accesible.

Más recursos

También puede consultar el repositorio llama-recipes y la guía Getting started with Llama sobre cómo realizar el ajuste fino utilizando FSDP y PEFT.

Precauciones

La fusión al usar PEFT y FSDP no se admite actualmente y generará un error.
Pasar el parámetro de configuración modules_to_save no se ha probado hasta el momento.
El ahorro de memoria de GPU al usar CPU Offloading no se ha probado hasta la fecha.
Al usar FSDP+QLoRA, paged_adamw_8bit actualmente provoca un error al guardar un punto de control.
El entrenamiento con DoRA y FSDP debería funcionar (aunque a una velocidad menor que con LoRA). Si se combina con bitsandbytes (QDoRA), la cuantización de 4 bits también debería funcionar, pero la cuantización de 8 bits tiene problemas conocidos y no se recomienda.
Con FSDP2, set_adapter() y set_requires_grad() redividen los módulos FSDP2 dentro de PeftModel, incluido el propio PeftModel, para que los cambios en requires_grad lleguen a los parámetros divididos. Al salir de disable_adapter(), solo se hace así si FSDP2 reemplazó los objetos de parámetros registrados en su interior, como lo hace una operación de propagación hacia adelante con los parámetros divididos, o si la reactivación cambia requires_grad respecto a su valor de entrada. Los módulos FSDP2 que dividen los parámetros de una capa de adaptador fusionado no se redividen, por lo que hay que desfusionar antes de llamar a estos métodos. Con reshard_after_forward=False, o para el módulo raíz con el valor predeterminado None, la siguiente propagación hacia adelante vuelve a realizar un all-gather. Los módulos FSDP2 que envuelven PeftModel no se redividen; llama a reshard() en ellos antes de estos métodos. Dado que estas llamadas pueden dejar los pesos base divididos, fusionar inmediatamente después no funciona correctamente (ver 1). Los métodos de aprendizaje de prompts y adaption prompt aún no están cubiertos.

Update on GitHub