Google verwendet KI-Technologie, um Inhalte in Ihre bevorzugte Sprache zu übersetzen. KI-Übersetzungen können Fehler enthalten.

Verstärkungslernen auf einer TPU-VM mit MaxText ausführen

In dieser detaillierten Anleitung erfahren Sie Schritt für Schritt, wie Sie das Training für bestärkendes Lernen (Reinforcement Learning, RL) auf einer einzelnen v6e-8 VM-Instanz (Virtual Machine) mit Tensor Processing Unit (TPU) Google Cloud mit MaxText ausführen. MaxText ist ein leistungsstarker JAX-basierter Trainingsstack für Large Language Models (LLMs).

Zusammenfassung der Arbeitslast

Für die Ausführung dieser RL-Nachbearbeitungs-Arbeitslast auf Cloud TPU sind die folgenden Komponenten erforderlich:

Modell:Llama 3.1 8B Instruct (llama3.1-8b-Instruct), ein LLM mit offenen Gewichten, das für Dialog- und anweisungsbezogene Aufgaben optimiert ist. In dieser Anleitung werden die Modellgewichte aus dem Hugging Face-Format in das MaxText-kompatible Format (mt-format) konvertiert und für die verteilte Ausführung auf einer v6e-8-TPU mit JAX vorbereitet.
Algorithmus für bestärkendes Lernen (Reinforcement Learning, RL):Group Relative Policy Optimization (GRPO). Dabei führt die Arbeitslast eine skalierbare Nachbearbeitungspipeline mit dem integrierten RL-Trainer von MaxText aus. Bei GRPO ist kein separates Wert- oder Kritikermodell erforderlich. Stattdessen wird eine Gruppe von Ausgaben bewertet, die für dieselbe Eingabeaufforderung generiert wurden. Dabei werden relative Belohnungen verwendet, um die Akteurrichtlinie effizient zu optimieren.
Dataset und Vorlage:Grade School Math 8K (GSM8K), ein hochwertiges Dataset mit sprachlich vielfältigen Textaufgaben für die Grundschule. Diese Arbeitslast verwendet es, um Rollouts zu generieren, Belohnungen zu berechnen und die mehrstufige Problemlösung während der RL-Abstimmungsphase zu verbessern.
Infrastruktur und Framework:Ein einzelner Cloud TPU v6e-8-Knoten (ct6e-standard-8t) mit MaxText (maxtext[tpu-post-train]), einem leistungsstarken Trainingsstack, der von JAX unterstützt wird.

Ziele

Cloud TPU-VM-Instanz einrichten
MaxText und seine Abhängigkeiten installieren
Hugging Face-Modell in das MaxText-Format konvertieren
RL-Arbeitslast (Group Relative Policy Optimization, GRPO) auf der TPU ausführen
Trainiertes Modell für die Bereitstellung wieder in das Hugging Face-Format konvertieren

Kosten

In diesem Dokument verwenden Sie die folgenden kostenpflichtigen Komponenten von Google Cloud:

Cloud TPU

Mit dem Preisrechner können Sie eine Kostenschätzung für Ihre voraussichtliche Nutzung vornehmen. Verwenden Sie den Preisrechner.

Neuen Google Cloud Nutzern vonsteht möglicherweise eine kostenlose Testversion zur Verfügung.

Nach Abschluss der in diesem Dokument beschriebenen Aufgaben können Sie weitere Kosten vermeiden, indem Sie die erstellten Ressourcen löschen. Weitere Informationen finden Sie unter Bereinigen.

Hinweis

Für diese Anleitung benötigen Sie ein Hugging Face-Zugriffstoken. Sie können sich kostenlos bei Hugging Face registrieren. Nachdem Sie ein Konto haben, generieren Sie ein Zugriffstoken:
1. Klicken Sie auf der Seite Welcome to Hugging Face (Willkommen bei Hugging Face) auf Ihren Kontoavatar und wählen Sie Access tokens (Zugriffstokens) aus.
2. Klicken Sie auf der Seite Access tokens (Zugriffstokens) auf Create new token (Neues Token erstellen).
3. Wählen Sie den Tokentyp Read (Lesen) aus und geben Sie einen Namen für Ihr Token ein.
4. Ihr Zugriffstoken wird angezeigt. Speichern Sie das Token an einem sicheren Ort.

Akzeptieren Sie auf der Hugging Face-Website die Lizenz vereinbarung für das Modell, das Sie trainieren möchten. In dieser Anleitung wird das Modell llama3.1-8b-Instruct verwendet.

Bitten Sie Ihren Administrator, Ihnen die folgenden IAM-Rollen für Ihr Projekt zuzuweisen, um die Berechtigungen zu erhalten, die Sie zum Ausführen dieser Anleitung benötigen:

TPU-Administrator (roles/tpu.admin)
Service Account User (roles/iam.serviceAccountUser)
Compute-Bearbeiter (roles/compute.editor)

Weitere Informationen zum Zuweisen von Rollen finden Sie unter Zugriff auf Projekte, Ordner und Organisationen verwalten.

Sie können die erforderlichen Berechtigungen auch über benutzerdefinierte Rollen oder andere vordefinierte Rollen erhalten.

Umgebung einrichten

Richten Sie die Umgebungsvariablen ein, indem Sie das folgende Skript ausführen:

export PROJECT="YOUR_PROJECT_ID"
export ZONE="YOUR_ZONE"
export RESERVATION="YOUR_RESERVATION_NAME"
export TPU_NAME="YOUR_TPU_NAME"

Ersetzen Sie Folgendes:

YOUR_PROJECT_ID: Ihre Google Cloud Projekt-ID
YOUR_ZONE: die Zone, die Sie verwenden möchten
YOUR_RESERVATION_NAME: Ihre Kapazitätsreservierung
YOUR_TPU_NAME: der Name Ihrer Cloud TPU-VM Instanz

Authentifizieren Sie sich mit Google Cloud indem Sie den folgenden Befehl ausführen:

gcloud auth login

Cloud TPU-VM erstellen

Erstellen Sie eine Cloud TPU-VM-Instanz mit 8 v6e-TPU-Chips, die an Ihre Kapazitätsreservierung gebunden ist.

gcloud compute instances create "${TPU_NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --machine-type=ct6e-standard-8t \
    --image-project=ubuntu-os-accelerator-images \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --boot-disk-size=200GB \
    --maintenance-policy=TERMINATE \
    --instance-termination-action=DELETE \
    --provisioning-model=RESERVATION_BOUND \
    --reservation-affinity=specific \
    --reservation="${RESERVATION}"

Nachdem die VM-Instanz erstellt wurde, stellen Sie über SSH eine Verbindung zu ihr her.

gcloud compute ssh "${TPU_NAME}" --zone "${ZONE}" --project "${PROJECT}"

Führen Sie die folgenden Schritte in Ihrer TPU-VM-Instanz aus.

MaxText installieren

Aktualisieren Sie die Systempakete in der TPU-VM-Instanz.

sudo apt update && sudo apt upgrade -y --fix-missing

Installieren Sie Python 3.12, das für MaxText erforderlich ist, und das zugehörige Paket für die virtuelle Umgebung.

sudo apt install -y build-essential cmake ninja-build

Verwenden Sie uv, um die Installation des Python-Pakets zu beschleunigen.

curl -LsSf https://astral.sh/uv/install.sh | sh
source "${HOME}/.local/bin/env"

Erstellen Sie eine virtuelle Umgebung mit dem Namen maxtext_venv und aktivieren Sie sie.

uv python install 3.12
uv venv --python 3.12 --seed maxtext_venv
source maxtext_venv/bin/activate

Installieren Sie MaxText und die Abhängigkeiten, die für Nachbearbeitungsaufgaben erforderlich sind.

UV_TORCH_BACKEND=cpu uv pip install "maxtext[tpu-post-train]==0.2.3" --resolution=lowest

Installieren Sie die verbleibenden erforderlichen Abhängigkeiten mit dem folgenden Befehl:

install_tpu_post_train_extra_deps

Modell in das MaxText-Format konvertieren

Wenn Sie das Modell im MaxText-Format trainieren möchten, müssen Sie es aus dem Hugging Face-Format in das MaxText-Format konvertieren.

Geben Sie folgende Werte an:

Ihr Hugging Face-Zugriffstoken
Der Name des Modells, das Sie verwenden möchten
Das Verzeichnis, in dem Sie das Modell im MaxText-Format speichern möchten
Optionen für das Laden und Speichern

export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME='llama3.1-8b-Instruct'
export MODEL_CHECKPOINT_DIRECTORY=/dev/shm/$MODEL_NAME/mt-format/
export USE_PATHWAYS=0 # Set to 1 for Pathways, 0 for McJAX
export LAZY_LOAD_TENSORS=False # True to use lazy load, False to use eager load.

Ersetzen Sie YOUR_HF_TOKEN durch das zuvor erstellte Hugging Face-Zugriffstoken.

Führen Sie das folgende Skript aus, um das Modell aus dem Hugging Face-Format in das MaxText-Format zu konvertieren. Diese Konvertierung dauert etwa fünf Minuten.

python3 -m maxtext.checkpoint_conversion.to_maxtext \
    model_name="${MODEL_NAME?}" \
    hf_access_token="${HF_TOKEN?}" \
    base_output_directory="${MODEL_CHECKPOINT_DIRECTORY?}" \
    scan_layers=True \
    use_multimodal=False \
    hardware=cpu \
    skip_jax_distributed_system=true \
    checkpoint_storage_use_zarr3=$((1 - USE_PATHWAYS)) \
    checkpoint_storage_use_ocdbt=$((1 - USE_PATHWAYS)) \
    --lazy_load_tensors="${LAZY_LOAD_TENSORS?}"

Trainingsarbeitslast starten

Nach Abschluss der Konvertierung können Sie die RL-Arbeitslast starten.

Konfigurieren Sie die Trainingsparameter für die RL-Arbeitslast.

# -- MaxText configuration --
export BASE_OUTPUT_DIRECTORY=/dev/shm/$MODEL_NAME/post-train/
RUN_NAME=$(date +%Y-%m-%d-%H-%M-%S)
export RUN_NAME
export CHIPS_PER_VM=8
export NUM_BATCHES=50
export MAXTEXT_CKPT_PATH=$MODEL_CHECKPOINT_DIRECTORY/0/items
export TPU_ACCELERATOR_TYPE=v6e-8
export TPU_WORKER_ID=0
TPU_NAME=$(hostname)
export TPU_NAME
export TPU_SKIP_MDS_QUERY=1
export TPU_WORKER_HOSTNAMES=localhost
export TPU_TOPOLOGY=2x4
export TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
export TPU_HOST_BOUNDS=1,1,1

Starten Sie den Trainingsjob. Auf einer v6e-8-VM-Instanz dauert das etwa 10 Minuten.

python3 -m maxtext.trainers.post_train.rl.train_rl \
    model_name="${MODEL_NAME?}" \
    load_parameters_path="${MAXTEXT_CKPT_PATH?}" \
    run_name="${RUN_NAME?}" \
    base_output_directory="${BASE_OUTPUT_DIRECTORY?}" \
    chips_per_vm="${CHIPS_PER_VM?}" \
    num_batches="${NUM_BATCHES?}" \
    num_test_batches=10 \
    rollout_data_parallelism=1 \
    rollout_tensor_parallelism=-1

Trainiertes Modell wieder in das Hugging Face-Format konvertieren

Nach Abschluss der Trainingsarbeitslast konvertieren Sie das Modell wieder in das Hugging Face-Format.

Legen Sie die Pfade für den Export und die trainierten Parameter fest.

export HF_EXPORT=/dev/shm/$MODEL_NAME/hf-trained/
export HF_MODEL_NAME=llama3.1-8b
export POST_TRAIN_PATH=$BASE_OUTPUT_DIRECTORY/$RUN_NAME/checkpoints/actor/$NUM_BATCHES/model_params

Führen Sie die Konvertierung zurück in das Hugging Face-Format aus.

python3 -m maxtext.checkpoint_conversion.to_huggingface \
    model_name="${HF_MODEL_NAME?}" \
    load_parameters_path="${POST_TRAIN_PATH?}" \
    base_output_directory="${HF_EXPORT?}" \
    scan_layers=True \
    use_multimodal=False \
    weight_dtype=bfloat16

Nach Abschluss der Konvertierung ist Ihr abgestimmtes Modell, das in /dev/shm/$MODEL_NAME/hf-trained gespeichert ist, einsatzbereit. Da Sie beim Neustart der VM den Zugriff auf die Inhalte des Ordners /dev/shm verlieren, sollten Sie das abgestimmte Modell in einen nichtflüchtigen Speicher verschieben oder in den Hugging Face Hub hochladen.

Bereinigen

Löschen Sie die in dieser Anleitung erstellten Ressourcen, um zusätzliche Gebühren zu vermeiden.

TPU-VM-Instanz löschen

Löschen Sie Ihre Cloud TPU-VM-Instanz.

gcloud compute instances delete "${TPU_NAME}" --zone="${ZONE}" --project="${PROJECT}" --quiet

Nächste Schritte

Weitere Informationen zu Cloud TPU finden Sie unter Einführung in Cloud TPU.
Details zur Architektur und Konfiguration der v6e-8-TPU finden Sie unter TPU v6e.