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

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

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 zur 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.

Wichtig: Für diese Anleitung benötigen Sie Zugriff auf v6e-8-TPUs. In dieser Anleitung wird davon ausgegangen, dass Sie bereits eine Reservierung für diese TPUs in Ihrem Google Cloud Projekt haben, damit die Kapazität bei Bedarf verfügbar ist. Weitere Informationen zu den Optionen zum Reservieren von Cloud TPU-Kapazität für ein zukünftiges Datum und eine zukünftige Uhrzeit finden Sie unter Cloud TPU-Reservierungen.

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 Ihre Umgebungsvariablen ein, indem Sie das folgende Skript ausführen:

export PROJECT="YOUR_PROJECT_ID"
export ZONE="ZONE_NAME"
export RESERVATION="RESERVATION_NAME"
export TPU_NAME="TPU_MACHINE_NAME"

Ersetzen Sie Folgendes:

• YOUR_PROJECT_ID: Ihre Google Cloud Projekt-ID
• ZONE_NAME: die Zone, die Sie verwenden möchten
• RESERVATION_NAME: Ihre Kapazitätsreservierung
• TPU_MACHINE_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 sind.

gcloud alpha compute tpus tpu-vm create $TPU_NAME \
    --zone=$ZONE \
    --project=$PROJECT \
    --accelerator-type=v6e-8 \
    --version=v2-alpha-tpuv6e \
    --provisioning-model=reservation-bound \
    --reservation=$RESERVATION

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

gcloud compute tpus tpu-vm 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 python3.12 python3.12-venv

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 venv --python 3.12 --seed maxtext_venv
source maxtext_venv/bin/activate

Installieren Sie MaxText und die Abhängigkeiten, die für Aufgaben nach dem Training erforderlich sind.

uv pip install maxtext[tpu-post-train]==0.2.2 --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 vom 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 vom 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.

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

   # -- MaxText configuration --
   export BASE_OUTPUT_DIRECTORY=/dev/shm/$MODEL_NAME/post-train/
   export RUN_NAME=$(date +%Y-%m-%d-%H-%M-%S)
   export CHIPS_PER_VM=8
   export NUM_BATCHES=50
   export MAXTEXT_CKPT_PATH=$MODEL_CHECKPOINT_DIRECTORY/0/items

2. 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.

1. 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

2. 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 unter /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 persistenten Speicher verschieben oder in den Hugging Face Hub hochladen.

Bereinigen

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

TPU-VM-Instanz löschen

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

gcloud alpha compute tpus tpu-vm delete $TPU_NAME --zone=$ZONE --project=$PROJECT --quiet

Nächste Schritte

• Weitere Informationen zu Cloud TPU finden Sie unter Einführung in Cloud TPU.
• Architektur- und Konfigurationsdetails für die v6e-8 TPU finden Sie unter TPU v6e.