HyperPod Slurm 集群 DPO 教程 (GPU) - 亚马逊 SageMaker AI
设置 HyperPod GPU Slurm 环境启动训练作业
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异，请参阅 中国的 Amazon Web Services 服务入门 (PDF)。

本文属于机器翻译版本。若本译文内容与英语原文存在差异，则一律以英文原文为准。

HyperPod Slurm 集群 DPO 教程 (GPU)

按照以下教程操作，设置 Slurm 环境并在 Llama 80 亿参数模型上开始直接偏好优化（DPO）。

先决条件

在开始设置环境之前，请确保您：

  • 设置 HyperPod GPU Slurm 集群

    • 你的 HyperPod Slurm 集群必须启用 Nvidia Enroot 和 Pyxis（默认情况下它们处于启用状态）。

  • 拥有一个共享存储位置。它可以是可从群集节点访问的 Amazon FSx 文件系统或 NFS 系统。

  • 采用下列格式之一的令牌化二元偏好数据集：

    • JSON

    • JSONGZ（压缩 JSON）

    • ARROW

  • （可选）如果您需要预先训练的权重， HuggingFace 或者您正在训练 Llama 3.2 模型，则必须在开始训练之前获取 HuggingFace 代币。有关获取令牌的更多信息，请参阅用户访问令牌。

设置 HyperPod GPU Slurm 环境

要在 Slurm 集群上启动训练作业，请执行以下操作：

  • 通过 SSH 连接到 Slurm 集群的头节点。

  • 登录后，设置虚拟环境。确保您使用 Python 3.9 或更高版本。

    #set up a virtual environment
python3 -m venv ${PWD}/venv
source venv/bin/activate

  • 将 SageMaker HyperPod 配方和 SageMaker HyperPod 适配器存储库克隆到共享存储位置。共享存储位置可以是可从群集节点访问的 Amazon FSx 文件系统或 NFS 系统。

    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

  • 使用 Enroot 创建一个 squash 文件。要查找最新版本的 SMP 容器，请参阅 SageMaker 模型并行性库发布说明。有关使用 Enroot 文件的更多信息，请参阅构建Amazon经过优化的 Nemo- Launcher 镜像。

    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>"

  • 要使用 Enroot squash 文件开始训练，请使用以下示例来修改 recipes_collection/config.yaml 文件。

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

启动训练作业

要在单个 Slurm 计算节点上为序列长度为 8192 的 Llama 80 亿参数模型启动 DPO 作业，请将启动脚本 launcher_scripts/llama/run_hf_llama3_8b_seq8k_gpu_dpo.sh 设置为以下内容：

  • IMAGE：环境设置部分中的容器。

  • HF_MODEL_NAME_OR_PATH：在配方的 hf_model_name_or_path 参数中定义预训练权重的名称或路径。

  • （可选）如果您需要预先训练的权重，则 HuggingFace 可以通过设置以下键值对来提供 HuggingFace 标记：

    recipes.model.hf_access_token=${HF_ACCESS_TOKEN}
注意

在此设置中，用于 DPO 的参考模型是从正在训练的基础模型自动派生出来的（未明确定义单独的参考模型）。DPO 特定的超参数已预配置了以下默认值：

  • beta：0.1（控制 KL 散度正则化的强度）

  • label_smoothing：0.0（未对偏好标签应用平滑处理）

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}$" \

在上一个脚本中配置所有必需的参数后，可以通过运行该脚本来启动训练作业。

bash launcher_scripts/llama/run_hf_llama3_8b_seq8k_gpu_dpo.sh

有关 Slurm 集群配置的更多信息，请参阅在 HyperPod Slurm 上运行训练作业。