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

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

Trainium Slurm 集群预训练教程

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

先决条件

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

  • 设置 SageMaker HyperPod Trainium Slurm 集群。

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

  • 以下格式之一的数据：

    • JSON

    • JSONGZ（压缩 JSON）

    • 箭头

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

在 Slurm 集群上设置 Trainium 环境

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

  • 通过 SSH 进入你的 Slurm 集群的主节点。

  • 登录后，设置 Neuron 环境。有关设置 Neuron 的信息，请参阅 Ne uron 设置步骤。我们建议依靠预装了 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 Llama 3-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 上运行训练作业