Trainium Slurm 集群预训练教程 - Amazon SageMaker AI
在 Slurm 集群上设置 Trainium 环境在 Trainium 中启动训练作业

Trainium Slurm 集群预训练教程

按照以下教程操作，在 Slurm 集群上设置 Trainium 环境，并在 Llama 80 亿参数模型上开始训练作业。

先决条件

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

  • 设置 SageMaker HyperPod Trainium Slurm 集群。

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

  • 拥有采用以下格式之一的数据：

    • JSON

    • JSONGZ（压缩 JSON）

    • ARROW

  • （可选）如果您使用 HuggingFace 中的模型权重进行预训练或微调，则必须获得 HuggingFace 令牌。有关获取令牌的更多信息，请参阅用户访问令牌。

在 Slurm 集群上设置 Trainium 环境

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

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

  • 登录后，设置 Neuron 环境。有关设置 Neuron 的信息，请参阅 Neuron 设置步骤。我们建议使用预安装了 Neuron 驱动程序的深度学习 AMI，例如带 DLAMI Pytorch 的 Ubuntu 20。

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

    git clone --recursive https://github.com/aws/sagemaker-hyperpod-recipes.git
cd sagemaker-hyperpod-recipes
pip3 install -r requirements.txt

  • 按照以下教程操作：HuggingFace Llama3-8B 预训练

  • 准备模型配置。Neuron 存储库中提供的模型配置。有关本教程中使用的模型配置，请参阅 llama3 8b 模型配置

在 Trainium 中启动训练作业

要在 Trainium 中启动训练作业，请指定集群配置和 Neuron 配方。例如，要在 Trainium 中启动 llama3 8b 预训练作业，请将启动脚本 launcher_scripts/llama/run_hf_llama3_8b_seq8k_trn1x4_pretrain.sh 设置为以下内容：

  • MODEL_CONFIG：环境设置部分中的模型配置

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

    recipes.model.hf_access_token=<your_hf_token>
#!/bin/bash

#Users should set up their cluster type in /recipes_collection/config.yaml

SAGEMAKER_TRAINING_LAUNCHER_DIR=${SAGEMAKER_TRAINING_LAUNCHER_DIR:-"$(pwd)"}

COMPILE=0
TRAIN_DIR="${TRAIN_DIR}" # Location of training dataset
MODEL_CONFIG="${MODEL_CONFIG}" # Location of config.json for the model

HYDRA_FULL_ERROR=1 python3 "${SAGEMAKER_TRAINING_LAUNCHER_DIR}/main.py" \
    base_results_dir="${SAGEMAKER_TRAINING_LAUNCHER_DIR}/results" \
    instance_type="trn1.32xlarge" \
    recipes.run.compile="$COMPILE" \
    recipes.run.name="hf-llama3-8b" \
    recipes.trainer.num_nodes=4 \
    recipes=training/llama/hf_llama3_8b_seq8k_trn1x4_pretrain \
    recipes.data.train_dir="$TRAIN_DIR" \
    recipes.model.model_config="$MODEL_CONFIG"

要启动训练作业，请运行以下命令：

bash launcher_scripts/llama/run_hf_llama3_8b_seq8k_trn1x4_pretrain.sh

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