Trainium SageMaker 训练作业预训练教程 - Amazon SageMaker AI
为 Trainium SageMaker 训练作业设置环境使用 Jupyter Notebook 启动训练作业使用配方启动程序启动训练作业

Trainium SageMaker 训练作业预训练教程

本教程将指导您完成使用 SageMaker 训练作业与 AWS Trainium 实例设置并运行预训练作业的过程。

  • 设置环境

  • 启动训练作业

在开始之前，确保您满足以下先决条件。

先决条件

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

  • Amazon FSx 文件系统或 S3 存储桶，可在其中加载数据和输出训练构件。

  • 在 Amazon SageMaker AI 请求 ml.trn1.32xlarge 实例的服务配额。要请求增加服务配额，请执行以下操作：

    请求增加 ml.trn1.32xlarge 实例的服务配额

    1. 导航到 AWS 服务配额控制台。

    2. 选择 AWS 服务。

    3. 选择 JupyterLab。

    4. 为 ml.trn1.32xlarge 指定一个实例。

  • 使用 AmazonSageMakerFullAccess 和 AmazonEC2FullAccess 托管策略创建一个 AWS Identity and Access Management（IAM）角色。这些策略可向 Amazon SageMaker AI 提供运行示例的权限。

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

    • JSON

    • JSONGZ（压缩 JSON）

    • ARROW

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

为 Trainium SageMaker 训练作业设置环境

在运行 SageMaker 训练作业之前，请使用 aws configure 命令来配置 AWS 凭证和首选区域。作为替代方案，您也可以通过环境变量（例如 AWS_ACCESS_KEY_ID、AWS_SECRET_ACCESS_KEY 和 AWS_SESSION_TOKEN）提供凭证。有关更多信息，请参阅 SageMaker AI Python SDK。

我们强烈建议在 SageMaker AI JupyterLab 中使用 SageMaker AI Jupyter Notebook 启动 SageMaker 训练作业。有关更多信息，请参阅SageMaker JupyterLab。

  • （可选）如果您在 Amazon SageMaker Studio 中使用 Jupyter Notebook，则可跳过运行以下命令。请务必使用 python 3.9 或更高版本

    # set up a virtual environment
python3 -m venv ${PWD}/venv
source venv/bin/activate
# install dependencies after git clone.

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

  • 安装 SageMaker AI Python SDK

    pip3 install --upgrade sagemaker

    • 如果您正在运行 llama 3.2 3.2 多模态训练作业，则 transformers 版本必须为 4.45.2 版或更高版本。

      • 仅在使用 SageMaker AI Python SDK 时将 transformers==4.45.2 追加到 source_dir 中的 requirements.txt。

      • 如果您正在使用 HyperPod 配方通过将 sm_jobs 用作集群类型来启动，则无需指定转换器版本。

    • Container：Neuron 容器由 SageMaker AI Python SDK 自动设置。

使用 Jupyter Notebook 启动训练作业

您可以使用以下 Python 代码通过配方运行 SageMaker 训练作业。它利用 SageMaker AI Python SDK 中的 PyTorch 估算器来提交配方。以下示例将 llama3-8b 配方作为 SageMaker AI 训练作业启动。

  • compiler_cache_url：用于保存已编译的构件（例如 Amazon S3 构件）的缓存。

import os
import sagemaker,boto3
from sagemaker.debugger import TensorBoardOutputConfig

from sagemaker.pytorch import PyTorch

sagemaker_session = sagemaker.Session()
role = sagemaker.get_execution_role()

recipe_overrides = {
    "run": {
        "results_dir": "/opt/ml/model",
    },
    "exp_manager": {
        "explicit_log_dir": "/opt/ml/output/tensorboard",
    },
    "data": {
        "train_dir": "/opt/ml/input/data/train",
    },
    "model": {
        "model_config": "/opt/ml/input/data/train/config.json",
    },
    "compiler_cache_url": "<compiler_cache_url>"
} 

tensorboard_output_config = TensorBoardOutputConfig(
    s3_output_path=os.path.join(output, 'tensorboard'),
    container_local_output_path=overrides["exp_manager"]["explicit_log_dir"]
)

estimator = PyTorch(
    output_path=output_path,
    base_job_name=f"llama-trn",
    role=role,
    instance_type="ml.trn1.32xlarge",
    sagemaker_session=sagemaker_session,
    training_recipe="training/llama/hf_llama3_70b_seq8k_trn1x16_pretrain",
    recipe_overrides=recipe_overrides,
)

estimator.fit(inputs={"train": "your-inputs"}, wait=True)

上一代码使用训练配方创建一个 PyTorch 估算器对象，然后使用 fit() 方法对模型进行拟合。使用 training_recipe 参数指定要用于训练的配方。

使用配方启动程序启动训练作业

  • 更新 ./recipes_collection/cluster/sm_jobs.yaml

    • compiler_cache_url：用于保存构件的 URL。它可以是 Amazon S3 URL。

    sm_jobs_config:
  output_path: <s3_output_path>
  wait: True
  tensorboard_config:
    output_path: <s3_output_path>
    container_logs_path: /opt/ml/output/tensorboard  # Path to logs on the container
  wait: True  # Whether to wait for training job to finish
  inputs:  # Inputs to call fit with. Set either s3 or file_system, not both.
    s3:  # Dictionary of channel names and s3 URIs. For GPUs, use channels for train and validation.
      train: <s3_train_data_path>
      val: null
  additional_estimator_kwargs:  # All other additional args to pass to estimator. Must be int, float or string.
    max_run: 180000
    image_uri: <your_image_uri>
    enable_remote_debug: True
    py_version: py39
  recipe_overrides:
    model:
      exp_manager:
        exp_dir: <exp_dir>
      data:
        train_dir: /opt/ml/input/data/train
        val_dir: /opt/ml/input/data/val

  • 更新 ./recipes_collection/config.yaml

    defaults:
  - _self_
  - cluster: sm_jobs
  - recipes: training/llama/hf_llama3_8b_seq8k_trn1x4_pretrain
cluster_type: sm_jobs # bcm, bcp, k8s or sm_jobs. If bcm, k8s or sm_jobs, it must match - cluster above.

instance_type: ml.trn1.32xlarge
base_results_dir: ~/sm_job/hf_llama3_8B # Location to store the results, checkpoints and logs.

  • 使用 main.py 启动作业

    python3 main.py --config-path recipes_collection --config-name config

有关配置 SageMaker 训练作业的更多信息，请参阅SageMaker 训练作业预训练教程（GPU）。