Trainium SageMaker 训练作业预训练教程 - 亚马逊 SageMaker AI
为 Trainium 培训工作 SageMaker 设置环境使用 Jupyter 笔记本启动训练作业使用食谱启动器启动训练作业
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异，请参阅 中国的 Amazon Web Services 服务入门 (PDF)。

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

Trainium SageMaker 训练作业预训练教程

本教程将指导您完成使用带有 Amazon Trainium 实例的训练作业设置和运行预 SageMaker 训练作业的过程。

  • 设置您的环境

  • 启动训练作业

在开始之前，请确保您具备以下先决条件。

先决条件

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

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

  • 在 Amazon A SageMaker I 上为该ml.trn1.32xlarge实例申请服务配额。要申请增加服务配额，请执行以下操作：

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

    1. 导航到 S Amazon ervice Quotas 控制台。

    2. 选择 Amazon 服务。

    3. 选择 JupyterLab。

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

  • 使用AmazonSageMakerFullAccess和AmazonEC2FullAccess托管策略创建 Amazon Identity and Access Management (IAM) 角色。这些政策为 Amazon SageMaker AI 提供了运行示例的权限。

  • 以下格式之一的数据：

    • JSON

    • JSONGZ（压缩 JSON）

    • 箭头

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

为 Trainium 培训工作 SageMaker 设置环境

在运行 SageMaker 训练作业之前，请使用aws configure命令配置您的 Amazon 凭据和首选区域。或者，您也可以通过环境变量（例如AWS_ACCESS_KEY_IDAWS_SECRET_ACCESS_KEY、和）提供证书AWS_SESSION_TOKEN。有关更多信息，请参阅 SageMaker AI Python 开发工具包。

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

  • （可选）如果您在 Amazon SageMaker Studio 中使用 Jupyter 笔记本，则可以跳过运行以下命令。一定要使用版本 >= 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 多模态训练作业，则transformers版本必须等于4.45.2或更高。

      • 仅当你使用 AI SageMaker Python SDK 时，才会追加transformers==4.45.2到 source_dir requirements.txt 中。

      • 如果您使用 HyperPod 配方sm_jobs作为集群类型启动，则无需指定变形金刚版本。

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

使用 Jupyter 笔记本启动训练作业

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

  • 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：用于保存工件的网址。它可以是亚马逊 S3 网址。

    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)。