Didacticiel de pré-entraînement des tâches d’entraînement Trainium SageMaker - Amazon SageMaker AI
Configuration de votre environnement pour les tâches d’entraînement Trainium SageMakerLancement de la tâche d’entraînement à l’aide d’un bloc-notes JupyterLancement de la tâche d’entraînement avec le lanceur de recettes

Didacticiel de pré-entraînement des tâches d’entraînement Trainium SageMaker

Ce didacticiel vous guide tout au long du processus de configuration et d’exécution d’une tâche de pré-entraînement à l’aide de tâches d’entraînement SageMaker avec des instances AWS Trainium.

  • Configuration de votre environnement

  • Lancement d’une tâche d’entraînement

Avant de commencer, assurez-vous d’avoir rempli les conditions préalables suivantes.

Prérequis

Avant de commencer à configurer votre environnement, assurez-vous d’avoir :

  • Le système de fichiers Amazon FSx ou un compartiment S3 dans lequel vous pouvez charger les données et générer les artefacts d’entraînement.

  • Demandé un quota de service pour l’instance ml.trn1.32xlarge sur Amazon SageMaker AI. Pour demander une augmentation du quota de service, procédez comme suit :

    Pour demander une augmentation du quota de service pour une instance ml.trn1.32xlarge

    1. Accédez à la console AWS Service Quotas.

    2. Choisissez Services AWS.

    3. Sélectionnez JupyterLab.

    4. Spécifiez une instance pour ml.trn1.32xlarge.

  • Créez un rôle Gestion des identités et des accès AWS (IAM) avec les politiques gérées AmazonSageMakerFullAccess et AmazonEC2FullAccess. Ces politiques fournissent à Amazon SageMaker AI les autorisations nécessaires pour exécuter les exemples.

  • Des données dans l’un des formats suivants :

    • JSON

    • JSONGZ (JSON compressé)

    • ARROW

  • (Facultatif) Si vous avez besoin des poids pré-entraînés de HuggingFace ou si vous entraînez un modèle Llama 3.2, vous devez obtenir le jeton HuggingFace avant de commencer l’entraînement. Pour plus d’informations sur l’obtention du jeton, consultez User access tokens.

Configuration de votre environnement pour les tâches d’entraînement Trainium SageMaker

Avant d’exécuter une tâche d’entraînement SageMaker, utilisez la commande aws configure pour configurer vos informations d’identification et votre région préférée AWS. Vous pouvez également fournir vos informations d’identification par le biais de variables d’environnement telles que AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY et AWS_SESSION_TOKEN. Pour plus d’informations, consultez SageMaker AI Python SDK.

Nous vous recommandons vivement d’utiliser un bloc-notes Jupyter SageMaker AI dans SageMaker AI JupyterLab pour lancer une tâche d’entraînement SageMaker. Pour plus d’informations, consultez SageMaker JupyterLab.

  • (Facultatif) Si vous utilisez un bloc-notes Jupyter dans Amazon SageMaker Studio, vous pouvez ignorer l’exécution de la commande suivante. Assurez-vous d’utiliser une version >= 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

  • Installez le kit SageMaker AI Python SDK.

    pip3 install --upgrade sagemaker

    • Si vous exécutez une tâche d’entraînement multimodale Llama 3.2, la version des transformers doit être 4.45.2 ou supérieure.

      • Ajoutez transformers==4.45.2 à requirements.txt dans source_dir uniquement lorsque vous utilisez le kit SageMaker AI Python SDK.

      • Si vous utilisez des formules HyperPod pour démarrer en utilisant sm_jobs comme type de cluster, vous n’avez pas besoin de spécifier la version des transformeurs.

    • Container : le conteneur Neuron est défini automatiquement par le kit SageMaker AI Python SDK.

Lancement de la tâche d’entraînement à l’aide d’un bloc-notes Jupyter

Vous pouvez utiliser le code Python suivant pour exécuter une tâche d’entraînement SageMaker en utilisant votre recette. Il utilise l’estimateur PyTorch issu du kit SageMaker AI Python SDK pour soumettre la recette. L’exemple suivant lance la recette llama3-8b en tant que tâche d’entraînement SageMaker AI.

  • compiler_cache_url : cache à utiliser pour enregistrer les artefacts compilés, tels qu’un artefact 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)

Le code précédent crée un objet estimateur PyTorch avec la recette d’entraînement, puis adapte le modèle à l’aide de la méthode fit(). Utilisez le paramètre training_recipe pour spécifier la recette que vous souhaitez utiliser pour l’entraînement.

Lancement de la tâche d’entraînement avec le lanceur de recettes

  • Mettre à jour ./recipes_collection/cluster/sm_jobs.yaml

    • compiler_cache_url : URL utilisée pour enregistrer les artefacts. Il peut s’agir d’une URL 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

  • Mettre à jour ./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.

  • Lancez la tâche avec main.py

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

Pour plus d’informations sur la configuration des tâches d’entraînement SageMaker, consultez Didacticiel de pré-entraînement des tâches d’entraînement SageMaker (GPU).