Tutorial zum Vortraining des Trainium-Slurm-Clusters - Amazon SageMaker AI
Einrichten der Trainium-Umgebung auf dem Slurm-ClusterStarten des Trainingsjobs in Trainium

Tutorial zum Vortraining des Trainium-Slurm-Clusters

Das folgende Tutorial richtet eine Trainium-Umgebung auf einem Slurm-Cluster ein und startet einen Trainingsjob auf einem Lama-Modell mit 8 Milliarden Parametern.

Voraussetzungen

Bevor Sie mit der Einrichtung Ihrer Umgebung beginnen, stellen Sie sicher, dass Sie über Folgendes verfügen:

  • Richten Sie einen Trainium-Slurm-Cluster von SageMaker HyperPod ein.

  • Ein gemeinsam genutzter Speicherort. Es kann sich um ein Amazon-FSx-Dateisystem oder NFS-System handeln, auf das von den Clusterknoten aus zugegriffen werden kann.

  • Daten in einem der folgenden Formate:

    • JSON

    • JSONGZ (komprimiertes JSON)

    • ARROW

  • (Optional) Sie benötigen ein HuggingFace-Token, wenn Sie die Modellgewichte von HuggingFace für das Training oder die Feinabstimmung verwenden. Weitere Informationen zum Abrufen des Tokens finden Sie unter Benutzerzugriffstoken.

Einrichten der Trainium-Umgebung auf dem Slurm-Cluster

Um einen Trainingsjob auf einem Slurm-Cluster zu initiieren, führen Sie die folgenden Schritte aus:

  • Verbinden Sie sich per SSH mit dem Hauptknoten Ihres Slurm-Clusters.

  • Nachdem Sie sich angemeldet haben, richten Sie die Neuron-Umgebung ein. Informationen zur Einrichtung von Neuron finden Sie unter Schritte zum Einrichten von Neuron. Wir empfehlen, sich auf die Deep-Learning-AMIs zu verlassen, die mit den Treibern von Neuron vorinstalliert sind, z. B. Ubuntu 20 mit DLAMI Pytorch

  • Klonen Sie das Rezept-Repository von SageMaker HyperPod auf einen gemeinsam genutzten Speicherort im Cluster. Beim gemeinsam genutzten Speicherort kann sich um ein Amazon-FSx-Dateisystem oder NFS-System handeln, auf das von den Clusterknoten aus zugegriffen werden kann.

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

  • Gehen Sie das folgende Tutorial durch: HuggingFace Llama3-8B Vortraining

  • Bereiten Sie eine Modellkonfiguration vor. Die im Neuron-Repo verfügbaren Modellkonfigurationen. Die in diesem Tutorial verwendete Modellkonfiguration finden Sie unter llama3 8b-Modellkonfiguration

Starten des Trainingsjobs in Trainium

Um einen Trainingsjob in Trainium zu starten, geben Sie eine Cluster-Konfiguration und ein Neuron-Rezept an. Um beispielsweise einen lama3 8b-Vortrainigsjob in Trainium zu starten, legen Sie das Startskript launcher_scripts/llama/run_hf_llama3_8b_seq8k_trn1x4_pretrain.sh wie folgt fest:

  • MODEL_CONFIG: Die Modellkonfiguration aus dem Abschnitt „Umgebungseinrichtung“

  • (Optional) Sie können das HuggingFace-Token bereitstellen, wenn Sie vortrainierte Gewichte von HuggingFace benötigen, indem Sie das folgende Schlüssel-Wert-Paar festlegen:

    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"

Um den Trainingsjob zu starten, führen Sie bitte den folgenden Befehl aus:

bash launcher_scripts/llama/run_hf_llama3_8b_seq8k_trn1x4_pretrain.sh

Weitere Informationen zur Konfiguration des Slurm-Clusters finden Sie unter Ausführen eines Trainingsjobs auf HyperPod Slurm.