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 diesem Tutorial wird Schritt für Schritt beschrieben, wie Sie das überwachte Fine-Tuning (Supervised Fine-Tuning, SFT) auf einer einzelnen v6e-8-TPU-VM-Instanz (Tensor Processing Unit, TPU) Google Cloud mit MaxText ausführen. MaxText ist ein leistungsstarker JAX-basierter Trainingsstack für Large Language Models (LLMs).

Ziele

Cloud TPU-VM-Instanz einrichten.
Installieren Sie MaxText und die zugehörigen Abhängigkeiten.
Konvertieren eines Hugging Face-Modells in das MaxText-Format
Führen Sie eine SFT-Trainingsarbeitslast auf der TPU aus.
Konvertieren Sie das feinabgestimmte Modell für die Bereitstellung zurück in das Hugging Face-Format.

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.

Neuen Nutzern von Google Cloud steht 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 dieses Tutorial benötigen Sie ein Hugging Face-Zugriffstoken. Sie können sich bei Hugging Face für ein kostenloses Konto registrieren. Nachdem Sie ein Konto erstellt haben, können Sie ein Zugriffstoken generieren:
1. Klicken Sie auf der Seite Welcome to Hugging Face (Willkommen bei Hugging Face) auf Ihren Konto-Avatar und wählen Sie Access tokens (Zugriffstokens) aus.
2. Klicken Sie auf der Seite Zugriffstokens auf Neues Token erstellen.
3. Wählen Sie den Tokentyp 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 Lizenzvereinbarung 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 bei 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 "${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 Ihre 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 ins MaxText-Format zu konvertieren. Die 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. Das dauert auf einer v6e-8-VM-Instanz 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

Nachdem der Trainings-Workload abgeschlossen ist, 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 kann das abgestimmte Modell, das in /dev/shm/gemma3-4b/hf-trained gespeichert ist, verwendet werden. 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 Gebühren zu vermeiden.

TPU-VM-Instanz löschen

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

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