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

In dieser Anleitung erfahren Sie Schritt für Schritt, wie Sie eine überwachte Feinabstimmung (Supervised Fine-Tuning, SFT) auf einer einzelnen v6e-8 VM-Instanz (Virtual Machine) mit Tensor Processing Unit (TPU) mit MaxText ausführen. MaxText ist ein leistungsstarker JAX-basierter Trainings-Stack für Large Language Models (LLMs). Google Cloud

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 feinabgestimmte Modell zur Bereitstellung wieder in das Hugging Face-Format konvertieren

Kosten

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

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 erstellt 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:

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 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 $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 $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 virtuelle Umgebungen.

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

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

    # -- MaxText configuration --
export BASE_OUTPUT_DIRECTORY=/dev/shm/$MODEL_NAME/post-train/
export RUN_NAME=$(date +%Y-%m-%d-%H-%M-%S)
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

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

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

  2. 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 feinabgestimmtes 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 feinabgestimmte 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 Ihre Cloud TPU-VM-Instanz und löschen Sie sie.

gcloud alpha compute tpus tpu-vm delete $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.