Ajuste fino de Llama 3 con FSDP y PEFT

Esta receta te guía a través del proceso de ajuste fino de un modelo Meta Llama 3 para la tarea de resumen de texto usando el conjunto de datos samsum en múltiples GPUs en un solo nodo o a través de varios nodos.

Requisitos

Asegúrate de haber instalado el paquete llama-cookbook (detalles).

También necesitaremos 2 paquetes:

1. PEFT para usar el ajuste fino eficiente en parámetros.
2. FSDP que nos ayuda a paralelizar el entrenamiento en múltiples GPUs. Más detalles.

[!NOTE] El paquete llama-cookbook instalará la versión 2.0.1 de PyTorch. En caso de que quieras usar FSDP con PEFT para el ajuste fino multi-GPU, por favor instala las versiones nightly de PyTorch (detalles)

La cuantificación INT8 no es compatible actualmente con FSDP

Cómo ejecutarlo

Obtén acceso a una máquina con múltiples GPUs (en este caso, probamos con 4 A100 y A10s).

Con FSDP + QLORA

Esto ha sido probado en 4 GPUs H100.

 FSDP_CPU_RAM_EFFICIENT_LOADING=1 ACCELERATE_USE_FSDP=1 torchrun --nnodes 1 --nproc_per_node 4  finetuning.py --enable_fsdp  --quantization int4 --model_name /path_of_model_folder/70B  --mixed_precision False --low_cpu_fsdp --use_peft --peft_method lora --output_dir Path/to/save/PEFT/model

Con FSDP + PEFT

torchrun --nnodes 1 --nproc_per_node 4  finetuning.py --enable_fsdp --model_name /path_of_model_folder/8B --use_peft --peft_method lora --output_dir Path/to/save/PEFT/model

# Change the num nodes and GPU per nodes in the script before running.
sbatch ./multi_node.slurm

Usamos torchrun para iniciar múltiples procesos para FSDP.

Los argumentos usados en el comando anterior son:

• --enable_fsdp bandera booleana para habilitar FSDP en el script
• --use_peft bandera booleana para habilitar los métodos PEFT en el script
• --peft_method para especificar el método PEFT, aquí usamos lora otras opciones son llama_adapter, prefix.

Solo con FSDP

Si te interesa ejecutar el ajuste fino de parámetros completos sin usar métodos PEFT, por favor usa el siguiente comando. Asegúrate de cambiar el nproc_per_node a tus GPUs disponibles. Esto ha sido probado con BF16 en 8xA100, GPUs de 40GB.

torchrun --nnodes 1 --nproc_per_node 8  finetuning.py --enable_fsdp --model_name /path_of_model_folder/8B --dist_checkpoint_root_folder model_checkpoints --dist_checkpoint_folder fine-tuned --fsdp_config.pure_bf16 --use_fast_kernels

Usando menos memoria de CPU (FSDP en modelo 70B)

Si estás ejecutando el ajuste fino de parámetros completos en el modelo 70B, puedes habilitar el modo low_cpu_fsdp como en el siguiente comando. Esta opción cargará el modelo solo en rank0 antes de mover el modelo a los dispositivos para construir FSDP. Esto puede ahorrar drásticamente la memoria de la CPU al cargar modelos grandes como el 70B (en un nodo de 8 GPUs, esto reduce la memoria de la CPU de más de 2T a 280G para el modelo 70B). Esto ha sido probado con BF16 en 16xA100, GPUs de 80GB.

torchrun --nnodes 1 --nproc_per_node 8 finetuning.py --enable_fsdp --low_cpu_fsdp --fsdp_config.pure_bf16 --model_name /path_of_model_folder/70B --batch_size_training 1 --dist_checkpoint_root_folder model_checkpoints --dist_checkpoint_folder fine-tuned

Multi-GPU multi-nodo:

Aquí usamos un script de slurm para programar un trabajo con slurm en múltiples nodos.

sbatch recipes/quickstart/finetuning/multi_node.slurm
# Change the num nodes and GPU per nodes in the script before running.

Para ajustar el modelo Meta Llama 405B con LoRA en 32xH100, GPUs de 80 GB, necesitamos combinar la cuantificación de 4 bits (QLoRA) y FSDP. Podemos lograr esto agregando las siguientes variables de entorno al script de slurm (antes del comando srun al final).

export FSDP_CPU_RAM_EFFICIENT_LOADING=1
export ACCELERATE_USE_FSDP=1 

Luego necesitamos reemplazar el comando srun inferior con el siguiente:

srun  torchrun --nproc_per_node 8 --rdzv_id $RANDOM --rdzv_backend c10d --rdzv_endpoint $head_node_ip:29500 ./finetuning.py  --enable_fsdp --use_peft --peft_method lora --quantization 4bit  --quantization_config.quant_type nf4 --mixed_precision False --low_cpu_fsdp

No olvides ajustar el número de nodos, ntasks y gpus-per-task en la parte superior.

Ejecutando con diferentes conjuntos de datos

Actualmente se admiten 3 conjuntos de datos de código abierto que se pueden encontrar en el archivo de configuración de conjuntos de datos. También puedes usar tu conjunto de datos personalizado (más información aquí).

• grammar_dataset : usa este notebook para obtener y procesar los conjuntos de datos Jfleg y C4 200M para la verificación gramatical.

• alpaca_dataset : para obtener estos datos de código abierto, por favor descarga el aplaca.json a la carpeta dataset.

wget -P ../../src/llama_cookbook/datasets https://raw.githubusercontent.com/tatsu-lab/stanford_alpaca/main/alpaca_data.json

• samsum_dataset

Para ejecutar con cada uno de los conjuntos de datos, establece la bandera dataset en el comando como se muestra a continuación:

# grammer_dataset
torchrun --nnodes 1 --nproc_per_node 4  finetuning.py --enable_fsdp  --model_name /path_of_model_folder/8B --use_peft --peft_method lora --dataset grammar_dataset --save_model --dist_checkpoint_root_folder model_checkpoints --dist_checkpoint_folder fine-tuned  --fsdp_config.pure_bf16 --output_dir Path/to/save/PEFT/model

# alpaca_dataset

torchrun --nnodes 1 --nproc_per_node 4  finetuning.py --enable_fsdp  --model_name /path_of_model_folder/8B --use_peft --peft_method lora --dataset alpaca_dataset --save_model --dist_checkpoint_root_folder model_checkpoints --dist_checkpoint_folder fine-tuned --fsdp_config.pure_bf16 --output_dir Path/to/save/PEFT/model


# samsum_dataset

torchrun --nnodes 1 --nproc_per_node 4  finetuning.py --enable_fsdp --model_name /path_of_model_folder/8B --use_peft --peft_method lora --dataset samsum_dataset --save_model --dist_checkpoint_root_folder model_checkpoints --dist_checkpoint_folder fine-tuned --fsdp_config.pure_bf16 --output_dir Path/to/save/PEFT/model

[CONSEJO] ¿Interconexión lenta entre nodos?

En caso de que estés lidiando con una red de interconexión más lenta entre nodos, para reducir la sobrecarga de comunicación puedes usar la bandera --hsdp.

HSDP (Hybrid sharding Data Parallel) ayuda a definir una estrategia de fragmentación híbrida donde puedes tener FSDP dentro de sharding_group_size, que puede ser el número mínimo de GPUs en las que puedes ajustar tu modelo, y DDP entre las réplicas del modelo especificadas por replica_group_size.

Esto requerirá establecer la estrategia de fragmentación en la configuración de fsdp en ShardingStrategy.HYBRID_SHARD y especificar dos configuraciones adicionales, sharding_group_size y replica_group_size, donde la primera especifica el tamaño del grupo de fragmentación, el número de GPUs en las que tu modelo puede caber para formar una réplica de un modelo, y la segunda especifica el tamaño del grupo de réplicas, que es world_size/sharding_group_size.

torchrun --nnodes 4 --nproc_per_node 8 ./finetuning.py --enable_fsdp --low_cpu_fsdp --fsdp_config.pure_bf16 --model_name /path_of_model_folder/70B --batch_size_training 1 --dist_checkpoint_root_folder model_checkpoints --dist_checkpoint_folder fine-tuned --hsdp --sharding_group_size n --replica_group_size world_size/n

Conteo de FLOPS y Perfilado de Pytorch

Para ayudar con el esfuerzo de benchmarking, estamos agregando soporte para contar los FLOPS durante el proceso de ajuste fino. Puedes lograr esto configurando --flop_counter al iniciar tu ajuste fino de una o varias GPUs. Usa --flop_counter_start para elegir qué paso contar los FLOPS. Se recomienda permitir una etapa de calentamiento antes de usar el contador de FLOPS.

De manera similar, puedes establecer la bandera --use_profiler y pasar una ruta de salida de perfilado usando --profiler_dir para capturar los rastros de perfil de tu modelo usando el perfilador de PyTorch. Para obtener un resultado de perfilado preciso, el perfilador de PyTorch requiere una etapa de calentamiento y la configuración actual es wait=1, warmup=2, active=3, por lo tanto, el perfilador comenzará el perfilado después del paso 3 y registrará los siguientes 3 pasos. Por lo tanto, para usar el perfilador de PyTorch, el --max-train-step debe ser mayor que 6. El perfilador de PyTorch sería útil para fines de depuración. Sin embargo, --flop_counter y --use_profiler no se pueden usar al mismo tiempo para garantizar la precisión de la medición.