Tutorial de DPO del clúster de HyperPod Slurm (GPU) - Amazon SageMaker AI

Tutorial de DPO del clúster de HyperPod Slurm (GPU)

En el siguiente tutorial se configura el entorno de Slurm y comienza un trabajo de optimización de preferencias directas (DPO, por sus siglas en inglés) en un modelo de 8000 millones de parámetros de Llama.

Requisitos previos

Antes de configurar el entorno, asegúrese de disponer de:

  • Configuración del clúster de HyperPod GPU Slurm

    • Su clúster de HyperPod Slurm debe tener activados Nvidia Enroot y Pyxis (están activados de forma predeterminada).

  • Una ubicación de almacenamiento compartida. Puede ser un sistema de archivos de Amazon FSx o un sistema NFS al que se pueda acceder desde los nodos del clúster.

  • Un conjunto de datos de preferencias binarias tokenizado en uno de los siguientes formatos:

    • JSON

    • JSONGZ (JSON comprimido)

    • ARROW

  • (Opcional) Si necesita las ponderaciones entrenadas previamente de HuggingFace o si está entrenando un modelo Llama 3.2, debe obtener el token de HuggingFace antes de empezar a entrenar. Para obtener más información sobre cómo obtener el token, consulte User access tokens.

Configuración del entorno de HyperPod GPU Slurm

Para iniciar un trabajo de entrenamiento en un clúster de Slurm, haga lo siguiente:

  • SSH en el nodo principal del clúster de Slurm.

  • Después de iniciar sesión, configure el entorno virtual. Asegúrese de utilizar Python 3.9 o posterior.

    #set up a virtual environment python3 -m venv ${PWD}/venv source venv/bin/activate
  • Clone las fórmulas de SageMaker HyperPod y los repositorios del adaptador de SageMaker HyperPod en una ubicación de almacenamiento compartida. La ubicación de almacenamiento compartido puede ser un sistema de archivos de Amazon FSx o un sistema NFS al que se pueda acceder desde los nodos del clúster.

    git clone https://github.com/aws/sagemaker-hyperpod-training-adapter-for-nemo.git git clone --recursive https://github.com/aws/sagemaker-hyperpod-recipes.git cd sagemaker-hyperpod-recipes pip3 install -r requirements.txt
  • Cree un archivo squash con Enroot. Para buscar la versión más reciente del contenedor de SMP, consulte Notas de la versión de la biblioteca de paralelismo de modelos de SageMaker. Para obtener más información sobre el uso del archivo Enroot, consulte Build AWS-optimized Nemo-Launcher image.

    REGION="<region>" IMAGE="658645717510.dkr.ecr.${REGION}.amazonaws.com/smdistributed-modelparallel:2.4.1-gpu-py311-cu121" aws ecr get-login-password --region ${REGION} | docker login --username AWS --password-stdin 658645717510.dkr.ecr.${REGION}.amazonaws.com enroot import -o $PWD/smdistributed-modelparallel.sqsh dockerd://${IMAGE} mv $PWD/smdistributed-modelparallel.sqsh "/fsx/<any-path-in-the-shared-filesystem>"
  • Para usar el archivo squash de Enroot para empezar a entrenar, utilice el siguiente ejemplo para modificar el archivo recipes_collection/config.yaml.

    container: /fsx/path/to/your/smdistributed-modelparallel.sqsh

Lanzamiento del trabajo de entrenamiento

Para lanzar un trabajo de DPO para el modelo de 8000 millones de parámetros de Llama con una longitud de secuencia de 8192 en un único nodo de computación de Slurm, defina el script de lanzamiento launcher_scripts/llama/run_hf_llama3_8b_seq8k_gpu_dpo.sh de la siguiente manera:

  • IMAGE: es el contenedor de la sección de configuración del entorno.

  • HF_MODEL_NAME_OR_PATH: defina el nombre o la ruta de las ponderaciones entrenadas previamente en el parámetro hf_model_name_or_path de la fórmula.

  • (Opcional) Puede proporcionar el token de HuggingFace si necesita ponderaciones entrenadas previamente de HuggingFace al configurar el siguiente par clave-valor:

    recipes.model.hf_access_token=${HF_ACCESS_TOKEN}
nota

El modelo de referencia utilizado para la DPO en esta configuración se deriva automáticamente del modelo base que se está entrenando (no se define explícitamente ningún modelo de referencia independiente). Los hiperparámetros específicos de la DPO se han configurado previamente con los siguientes valores predeterminados:

  • beta: 0,1 (controla la intensidad de la regularización de la divergencia KL)

  • label_smoothing: 0,0 (no se suaviza ninguna etiqueta de preferencias)

recipes.dpo.beta=${BETA} recipes.dpo.label_smoothing=${LABEL_SMOOTHING}
#!/bin/bash IMAGE="${YOUR_IMAGE}" SAGEMAKER_TRAINING_LAUNCHER_DIR="${SAGEMAKER_TRAINING_LAUNCHER_DIR:-${PWD}}" TRAIN_DIR="${YOUR_TRAIN_DIR}" # Location of training dataset VAL_DIR="${YOUR_VAL_DIR}" # Location of validation dataset # experiment output directory EXP_DIR="${YOUR_EXP_DIR}" HF_ACCESS_TOKEN="${YOUR_HF_TOKEN}" HF_MODEL_NAME_OR_PATH="${HF_MODEL_NAME_OR_PATH}" BETA="${BETA}" LABEL_SMOOTHING="${LABEL_SMOOTHING}" # Add hf_model_name_or_path and turn off synthetic_data HYDRA_FULL_ERROR=1 python3 ${SAGEMAKER_TRAINING_LAUNCHER_DIR}/main.py \ recipes=fine-tuning/llama/hf_llama3_8b_seq8k_gpu_dpo \ base_results_dir=${SAGEMAKER_TRAINING_LAUNCHER_DIR}/results \ recipes.run.name="hf_llama3_dpo" \ recipes.exp_manager.exp_dir="$EXP_DIR" \ recipes.model.data.train_dir="$TRAIN_DIR" \ recipes.model.data.val_dir="$VAL_DIR" \ recipes.model.hf_model_name_or_path="$HF_MODEL_NAME_OR_PATH" \ container="${IMAGE}" \ +cluster.container_mounts.0="/fsx:/fsx" \ recipes.model.hf_access_token="${HF_ACCESS_TOKEN}" \ recipes.dpo.enabled=true \ recipes.dpo.beta="${BETA}" \ recipes.dpo.label_smoothing="${LABEL_SMOOTHING}$" \

Tras configurar todos los parámetros necesarios en el script anterior, puede iniciar el trabajo de entrenamiento ejecutándolo.

bash launcher_scripts/llama/run_hf_llama3_8b_seq8k_gpu_dpo.sh

Para obtener más información acerca de la configuración del clúster de Slurm, consulte Ejecución de un trabajo de entrenamiento en HyperPod Slurm.