Exécuter l'affinage supervisé sur une VM TPU à l'aide de MaxText

Ce tutoriel fournit un guide par étapes pour exécuter un affinage supervisé (SFT) sur une seule instance de machine virtuelle (VM) v6e-8 Tensor Processing Unit (TPU) sur Google Cloud à l'aide de MaxText, une pile d'entraînement hautes performances basée sur JAX pour les grands modèles de langage (LLM).

Objectifs

• Configurez une instance de VM Cloud TPU.
• Installez MaxText et ses dépendances.
• Convertissez un modèle Hugging Face au format MaxText.
• Exécutez une charge de travail d'entraînement SFT sur le TPU.
• Reconvertissez le modèle affiné au format Hugging Face pour le diffuser.

Coûts

Dans ce document, vous utilisez les composants facturables suivants de Google Cloud :

• Cloud TPU

Pour obtenir une estimation des coûts en fonction de votre utilisation prévue, utilisez le simulateur de coût.

Les nouveaux utilisateurs de Google Cloud peuvent bénéficier d'un essai sans frais.

Une fois que vous avez terminé les tâches décrites dans ce document, supprimez les ressources que vous avez créées pour éviter que des frais vous soient facturés. Pour en savoir plus, consultez la section Effectuer un nettoyage.

Important : Pour suivre ce tutoriel, vous devez avoir accès à des TPU v6e-8. Ce tutoriel suppose que vous disposez déjà d'une réservation pour ces TPU dans votre projetGoogle Cloud afin de vous assurer que la capacité est disponible lorsque vous en avez besoin. Pour en savoir plus sur les options de réservation de capacité Cloud TPU pour une date et une heure ultérieures, consultez À propos des réservations Cloud TPU.

Avant de commencer

• Vous avez besoin d'un jeton d'accès Hugging Face pour suivre ce tutoriel. Vous pouvez créer un compte sans frais sur Hugging Face. Une fois que vous avez un compte, générez un jeton d'accès :

  1. Sur la page Bienvenue sur Hugging Face, cliquez sur l'avatar de votre compte, puis sélectionnez Jetons d'accès.
  2. Sur la page Jetons d'accès, cliquez sur Créer un jeton.
  3. Sélectionnez le type de jeton Lecture et saisissez un nom pour votre jeton.
  4. Votre jeton d'accès s'affiche. Enregistrez le jeton dans un endroit sûr.

• Sur le site Web Hugging Face, acceptez le contrat de licence du modèle que vous prévoyez d'entraîner. Ce tutoriel utilise le modèle gemma3-4b.

Pour obtenir les autorisations nécessaires pour suivre ce tutoriel, demandez à votre administrateur de vous accorder les rôles IAM suivants sur votre projet :

• Administrateur TPU (roles/tpu.admin)
• Utilisateur du compte de service (roles/iam.serviceAccountUser)
• Éditeur Compute (roles/compute.editor)

Pour en savoir plus sur l'attribution de rôles, consultez Gérer l'accès aux projets, aux dossiers et aux organisations.

Vous pouvez également obtenir les autorisations requises avec des rôles personnalisés ou d'autres rôles prédéfinis.

Configurer l'environnement

Configurez vos variables d'environnement en exécutant le script suivant :

export PROJECT="YOUR_PROJECT_ID"
export ZONE="YOUR_ZONE"
export RESERVATION="YOUR_RESERVATION_NAME"
export NAME="YOUR_TPU_NAME"
export NETWORK="default"

Remplacez les éléments suivants :

• YOUR_PROJECT_ID : ID de votre projet Google Cloud
• YOUR_ZONE : zone que vous souhaitez utiliser
• YOUR_RESERVATION_NAME : votre réservation de capacité
• YOUR_TPU_NAME : nom de votre instance de VM Cloud TPU

Authentifiez-vous auprès de Google Cloud en exécutant la commande suivante :

gcloud auth login

Créer votre VM Cloud TPU

Créez une instance de VM Cloud TPU avec huit puces TPU v6e, liées à votre réservation de capacité.

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}"

Une fois l'instance de VM créée, connectez-vous à celle-ci à l'aide de SSH.

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

Suivez les étapes ci-dessous dans votre instance de VM TPU.

Installer MaxText

Mettez à jour les packages système dans l'instance de VM TPU.

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

Installez Python 3.12, qui est requis par MaxText, ainsi que son package d'environnement virtuel.

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

Utilisez uv pour accélérer l'installation du package Python.

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

Créez un environnement virtuel nommé maxtext_venv et activez-le.

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

Installez MaxText et les dépendances requises pour les tâches post-entraînement.

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

Installez les dépendances requises restantes en exécutant la commande suivante :

install_tpu_post_train_extra_deps

Convertir le modèle au format MaxText

Pour entraîner le modèle au format MaxText, vous devez le convertir du format Hugging Face au format MaxText.

Spécifiez vos variables d'environnement, telles que votre jeton d'accès Hugging Face, le nom du modèle que vous souhaitez utiliser et le répertoire dans lequel vous souhaitez enregistrer le modèle au format MaxText.

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.

Remplacez YOUR_HF_TOKEN par le jeton d'accès Hugging Face que vous avez créé précédemment.

Pour convertir le modèle du format Hugging Face au format MaxText, exécutez le script suivant. Cette conversion prend environ cinq minutes.

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?}"

Démarrer la charge de travail d'entraînement

Une fois le processus de conversion terminé, vous pouvez démarrer la charge de travail SFT.

1. Configurez les paramètres d'entraînement de la charge de travail SFT.

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

2. Démarrez la tâche d'entraînement. Cela prend environ 10 minutes sur une instance de VM v6e-8.

   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

Reconvertir le modèle entraîné au format Hugging Face

Une fois la charge de travail d'entraînement terminée, reconvertissez le modèle au format Hugging Face.

1. Définissez les chemins d'exportation et les paramètres entraînés.

   export HF_EXPORT=/dev/shm/$MODEL_NAME/hf-trained/
   export POST_TRAIN_PATH=$BASE_OUTPUT_DIRECTORY/$RUN_NAME/checkpoints/$STEPS/model_params

2. Exécutez la conversion au format Hugging Face.

   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

Une fois la conversion terminée, votre modèle réglé stocké dans /dev/shm/gemma3-4b/hf-trained est prêt à être utilisé. Étant donné que vous perdez l'accès au contenu du dossier /dev/shm lorsque la VM redémarre, vous devez déplacer le modèle ajusté vers un stockage persistant ou l'importer dans le Hugging Face Hub.

Effectuer un nettoyage

Pour éviter des frais supplémentaires, supprimez les ressources créées au cours de ce tutoriel.

Supprimer votre instance de VM TPU

Quittez votre instance de VM Cloud TPU, puis supprimez-la.

# 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

Étapes suivantes

• Pour en savoir plus sur Cloud TPU, consultez la Présentation de Cloud TPU.
• Pour en savoir plus sur l'architecture et la configuration du TPU v6e-8, consultez TPU v6e.