Tutorial de pré-treinamento de tarefas de treinamento do SageMaker com o Trainium - Amazon SageMaker AI
Configurar o ambiente para tarefas de treinamento do SageMaker com o TrainiumIniciar a tarefa de treinamento com um caderno JupyterIniciar a tarefa de treinamento com o inicializador de fórmulas

Tutorial de pré-treinamento de tarefas de treinamento do SageMaker com o Trainium

Este tutorial orienta você pelo processo de configuração e execução de um trabalho de pré-treinamento usando tarefas de treinamento do SageMaker com instâncias do AWS Trainium.

  • Configure o ambiente.

  • Iniciar uma tarefa de treinamento

Antes de começar, verifique se você atendeu aos pré-requisitos a seguir.

Pré-requisitos

Antes de começar a configurar seu ambiente, você deve:

  • Sistema de arquivos do Amazon FSx ou o bucket S3 onde você pode carregar os dados e gerar os artefatos de treinamento.

  • Solicite uma cota de serviço para a instância ml.trn1.32xlarge no Amazon SageMaker AI. Para solicitar um aumento da cota de serviço, você pode seguir um dos seguintes procedimentos:

    Como solicitar um aumento de cota de serviço para a instância ml.trn1.32xlarge

    1. Navegue até o console do AWS Service Quotas.

    2. Selecione serviços da AWS.

    3. Escolha JupyterLab.

    4. Especifique 1 instância para ml.trn1.32xlarge.

  • Crie um perfil do AWS Identity and Access Management (IAM) com as políticas gerenciadas AmazonSageMakerFullAccess e AmazonEC2FullAccess. Essas políticas fornecem permissões ao Amazon SageMaker AI para executar os exemplos.

  • Ter dados em um dos seguintes formatos:

    • JSON

    • JSONGZ (JSON compactado)

    • ARROW

  • (Opcional) Se você precisar dos pesos pré-treinados do HuggingFace ou se estiver treinando um modelo Llama 3.2, deverá obter o token do HuggingFace antes de começar a treinar. Para ter mais informações sobre como obter o token, consulte User access tokens.

Configurar o ambiente para tarefas de treinamento do SageMaker com o Trainium

Antes de executar uma tarefa de treinamento do SageMaker, use o comando aws configure para configurar suas credenciais da AWS e a região de sua preferência. Também é possível fornecer suas credenciais por meio de variáveis de ambiente, como AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY e AWS_SESSION_TOKEN. Para ter mais informações, consulte SageMaker AI Python SDK.

É altamente recomendável usar um caderno Jupyter do SageMaker AI no SageMaker AI JupyterLab para iniciar uma tarefa de treinamento do SageMaker. Para ter mais informações, consulte SageMaker JupyterLab.

  • (Opcional) Se você estiver usando o caderno Jupyter no Amazon SageMaker Studio, poderá ignorar a execução do comando a seguir. Use uma versão >= à 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

  • Instale o SageMaker AI Python SDK.

    pip3 install --upgrade sagemaker

    • Se você estiver executando uma tarefa de treinamento multimodal do llama 3.2, a versão do transformers deve ser 4.45.2 ou posterior.

      • Acrescente transformers==4.45.2 a requirements.txt em source_dir somente quando estiver usando o SageMaker AI Python SDK.

      • Se você estiver usando fórmulas do HyperPod para iniciar com sm_jobs como o tipo de cluster, não precisará especificar a versão dos transformadores.

    • Container: o contêiner do Neuron é configurado automaticamente pelo SageMaker AI Python SDK.

Iniciar a tarefa de treinamento com um caderno Jupyter

Você pode usar o código Python a seguir para executar uma tarefa de treinamento do SageMaker usando sua fórmula. Ele usa o estimador do PyTorch do SageMaker AI Python SDK para enviar a fórmula. O exemplo a seguir inicia a fórmula do llama3-8b como uma tarefa de treinamento do SageMaker AI.

  • compiler_cache_url: cache a ser usado para salvar os artefatos compilados, como um artefato do 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)

O código anterior cria um objeto estimador do PyTorch com a fórmula de treinamento e, em seguida, ajusta o modelo usando o método fit(). Use o parâmetro training_recipe para especificar a fórmula que você deseja usar para treinamento.

Iniciar a tarefa de treinamento com o inicializador de fórmulas

  • Atualizar as ./recipes_collection/cluster/sm_jobs.yaml

    • compiler_cache_url: o URL usado para salvar os artefatos. Pode ser um URL do Amazon 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

  • Atualizar as ./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.

  • Iniciar a tarefa com main.py

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

Para ter mais informações sobre como configurar tarefas de treinamento do SageMaker, consulte Tutorial de pré-treinamento de tarefas de treinamento do SageMaker (GPU).