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

Überwachte Feinabstimmung mit MaxText auf einer TPU-VM ausführen

In dieser Anleitung finden Sie eine detaillierte Anleitung zum Ausführen der überwachten Feinabstimmung (Supervised Fine-Tuning, SFT) auf einer einzelnen v6e-8 VM-Instanz (Virtual Machine) der Tensor Processing Unit (TPU) mit MaxText, einem leistungsstarken JAX-basierten Trainings stack für Large Language Models (LLMs). Google Cloud

Zusammenfassung der Arbeitslast

Für die Ausführung dieser SFT-Arbeitslast auf einer Cloud TPU-VM sind die folgenden Komponenten erforderlich:

Modell:Gemma 3 4B (gemma3-4b), ein leichtgewichtiges Sprachmodell mit offenen Gewichten. In dieser Anleitung werden die Modellgewichte vom Hugging Face-Format in das MaxText-kompatible Format (mt-format) konvertiert und für die Ausführung auf einer einzelnen Cloud TPU v6e-8-VM mit JAX vorbereitet.
Trainingsmethode:SFT mit MaxText. Die Arbeitslast führt eine Fine-Tuning-Pipeline auf einer einzelnen Cloud TPU v6e-8-Instanz mit JAX aus. Dabei werden die Modellparameter in über 1.000 Schritten mit dem integrierten SFT-Trainer von MaxText (maxtext.trainers.post_train.sft.train_sft) und dialogorientierten Datasets für das Instruction-Tuning optimiert.
Dataset und Vorlage:UltraChat 200k (HuggingFaceH4/ultrachat_200k), ein umfangreiches, hochwertiges Dataset für Dialoge mit mehreren Turns (train_sft-Split mit der Spalte messages). Es wird verwendet, um das Modell zu optimieren und seine Fähigkeiten in Bezug auf Unterhaltungen und das Befolgen von Anweisungen zu verbessern.
Infrastruktur und Framework:eine einzelne Cloud TPU v6e-8-VM-Instanz (ct6e-standard-8t). Das Framework ist MaxText (maxtext[tpu-post-train]) mit JAX.

Ziele

Cloud TPU-VM-Instanz einrichten
MaxText und seine Abhängigkeiten installieren
Hugging Face-Modell in das MaxText-Format konvertieren
SFT-Trainingsarbeitslast auf der TPU ausführen
Das optimierte 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.

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 gemma3-4b 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 NAME="YOUR_TPU_NAME"
export NETWORK="default"

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

gcloud compute instances create "${NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --network="${NETWORK}" \
    --tags="${NAME}" \
    --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 "${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 Aufgaben nach dem Training erforderlich sind.

UV_TORCH_BACKEND=cpu 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 die Umgebungsvariablen an, z. B. Ihr Hugging Face-Zugriffstoken, den Namen des Modells, das Sie verwenden möchten, und das Verzeichnis, in dem Sie das Modell im MaxText-Format speichern möchten.

export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME='gemma3-4b'
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 SFT-Arbeitslast starten.

Konfigurieren Sie die Trainingsparameter für die SFT-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 STEPS=1000
export PER_DEVICE_BATCH_SIZE=1

# -- Dataset configuration --
export DATASET_NAME="HuggingFaceH4/ultrachat_200k"
export TRAIN_SPLIT="train_sft"
export TRAIN_DATA_COLUMNS="['messages']"

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.sft.train_sft \
    run_name="${RUN_NAME?}" \
    base_output_directory="${BASE_OUTPUT_DIRECTORY?}" \
    model_name="${MODEL_NAME?}" \
    load_parameters_path="${MAXTEXT_CKPT_PATH?}" \
    per_device_batch_size="${PER_DEVICE_BATCH_SIZE?}" \
    steps="${STEPS?}" \
    hf_path="${DATASET_NAME?}" \
    train_split="${TRAIN_SPLIT?}" \
    train_data_columns="${TRAIN_DATA_COLUMNS?}" \
    profiler=xplane

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 POST_TRAIN_PATH=$BASE_OUTPUT_DIRECTORY/$RUN_NAME/checkpoints/$STEPS/model_params

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

python3 -m maxtext.checkpoint_conversion.to_huggingface \
    model_name="${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 optimiertes Modell, das unter /dev/shm/gemma3-4b/hf-trained gespeichert ist, einsatzbereit. Da Sie beim Neustart der VM den Zugriff auf die Inhalte des Ordners /dev/shm verlieren, sollten Sie das optimierte 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 Kosten zu vermeiden.

TPU-VM-Instanz löschen

Beenden Sie die Cloud TPU-VM-Instanz und löschen Sie sie.

# 1. Delete TPU instance
echo "Deleting TPU instance: ${NAME}..."
gcloud compute instances delete "${NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --quiet || true

# 2. Delete IAP firewall rule
echo "Deleting firewall rule: ${FIREWALL_RULE_NAME}..."
gcloud compute firewall-rules delete "${FIREWALL_RULE_NAME}" \
    --project="${PROJECT}" \
    --quiet || true

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.