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

Die vorliegende Übersetzung wurde maschinell erstellt. Im Falle eines Konflikts oder eines Widerspruchs zwischen dieser übersetzten Fassung und der englischen Fassung (einschließlich infolge von Verzögerungen bei der Übersetzung) ist die englische Fassung maßgeblich.

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 SageMaker HyperPod Trainium Slurm-Cluster ein.

  • Ein gemeinsam genutzter Speicherort. Es kann sich um ein FSx Amazon-Dateisystem oder ein 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 SageMaker HyperPod Rezept-Repository an einen gemeinsam genutzten Speicherort im Cluster. Der gemeinsam genutzte Speicherort kann ein FSx Amazon-Dateisystem oder ein NFS-System sein, 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 Pretraining

  • 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 angeben, wenn Sie vorab trainierte Gewichte von benötigen, HuggingFace 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 Einen Trainingsjob auf HyperPod Slurm ausführen.