Tutorial de Fine-Tuning para modelos Llama4 con torchtune

Tutorial de Fine-Tuning para modelos Llama4 con torchtune

Este tutorial te muestra cómo realizar fine-tuning en modelos Llama4 usando torchtune.

Requisitos previos

0. Los modelos Llama 4 son grandes, por lo que te recomendamos usar 8 x H100 con 80 GB de RAM disponible para ejecutar el fine-tuning. Si usas GPUs con 40 GB de RAM, considera habilitar la descarga de memoria a la CPU. Si alquilas una VM de un proveedor de la nube para hacer el fine-tuning, asegúrate de tener los controladores NVIDIA instalados: wget https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run sudo sh cuda_12.1.0_530.30.02_linux.run

Verifica las GPUs con el comando nvidia-smi.

1. Necesitamos usar torchtune para realizar el fine-tuning de LoRA. Ahora, configuremos el entorno e instalemos la versión nightly de PyTorch. El fine-tuning de Llama4 LORA también requiere una compilación desde el código fuente.

a. Crea un entorno Python 3.10 (recomendado):

conda create -n py310 python=3.10 -y
conda activate py310

b. Instala PyTorch nightly y TorchTune:

pip install --force-reinstall --pre torch torchvision torchao --index-url https://download.pytorch.org/whl/nightly/cu126
git clone https://github.com/pytorch/torchtune.git
cd torchtune
git checkout 5d51c25cedfb6ba7b00e03cb2fef4f9cdb7baebd
pip install -e .

c. Dependiendo de la configuración de tu entorno, es posible que también necesites instalar estos paquetes:

pip install importlib_metadata
pip install torchvision
pip install torchao

2. También necesitamos un token de acceso de Hugging Face (HF_TOKEN) para la descarga de modelos; sigue las instrucciones aquí para obtener tu propio token. También necesitarás obtener acceso a los modelos Llama4 desde aquí

Pasos

1. Descarga los pesos de Llama4

Usaremos meta-llama/Llama-4-Scout-17B-16E-Instruct como ejemplo aquí. Reemplaza con tu token de Hugging Face:

tune download meta-llama/Llama-4-Scout-17B-16E-Instruct --output-dir /tmp/Llama-4-Scout-17B-16E-Instruct --hf-token $HF_TOKEN

Alternativamente, puedes usar huggingface-cli para iniciar sesión y luego descargar los pesos del modelo.

huggingface-cli login --token $HF_TOKEN
tune download meta-llama/Llama-4-Scout-17B-16E-Instruct --output-dir /tmp/Llama-4-Scout-17B-16E-Instruct

Esto recupera los pesos del modelo y el tokenizador de Hugging Face.

2. Ejecuta el Fine-Tuning de LoRA para Llama4

Para ejecutar el fine-tuning de LoRA, usa el siguiente comando:

tune run --nproc_per_node 8 lora_finetune_distributed --config llama4/scout_17B_16E_lora

Esto ejecutará el fine-tuning de LoRA en el modelo Llama4 con 8 GPUs. La configuración llama4/scout_17B_16E_lora es un archivo de configuración que especifica el modelo, el tokenizador y los parámetros de entrenamiento. Este comando también descargará el alpaca_dataset como se seleccionó en el archivo de configuración. Consulta la sección de Datasets para más detalles.

Puedes personalizar el proceso de entrenamiento agregando anulaciones de línea de comandos. Por ejemplo, para optimizar la eficiencia del entrenamiento y el uso de memoria:

  tune run --nproc_per_node 8 lora_finetune_distributed --config llama4/scout_17B_16E_lora batch_size=4 dataset.packed=True tokenizer.max_seq_len=2048 fsdp_cpu_offload=True

Estos argumentos controlan:

• batch_size=4: Establece el número de ejemplos procesados en cada paso de entrenamiento.
• dataset.packed=True: Empaqueta múltiples secuencias en un solo ejemplo para maximizar la utilización de la GPU.
• tokenizer.max_seq_len=2048: Establece la longitud máxima de la secuencia en 2048 tokens.
• fsdp_cpu_offload=True: Habilita la descarga de memoria a la CPU para evitar errores de falta de memoria (OOM), pero ralentiza significativamente el proceso de fine-tuning (especialmente importante si usas GPUs con 40 GB).

Consulta este archivo yaml para ver todas las configuraciones posibles para anular. Para obtener más información sobre la configuración YAML, consulta la documentación de configuración YAML

3. Ejecuta el Fine-Tuning de Parámetros Completos para Llama4

Para ejecutar el fine-tuning de parámetros completos, usa el siguiente comando:

  tune run  --nproc_per_node 8 full_finetune_distributed --config llama4/scout_17B_16E_full batch_size=4 dataset.packed=True tokenizer.max_seq_len=2048

Este comando ejecutará un fine-tuning completo en un solo nodo, ya que Torchtune, por defecto, usa la descarga a la CPU para evitar errores de falta de memoria (OOM). Consulta este archivo yaml para ver todas las configuraciones posibles para anular.

Alternativamente, si quieres ejecutarlo con múltiples nodos para evitar una posible lentitud por la descarga a la CPU, modifica este script de slurm.