Esegui il fine-tuning supervisionato su una VM TPU utilizzando MaxText

Questo tutorial fornisce una guida passo passo per l'esecuzione del fine tuning supervisionato (SFT) su una singola istanza di macchina virtuale (VM) di v6e-8Tensor Processing Unit (TPU) su Google Cloud utilizzando MaxText, uno stack di addestramento basato su JAX ad alte prestazioni per modelli linguistici di grandi dimensioni (LLM).

Obiettivi

• Configura un'istanza VM Cloud TPU.
• Installa MaxText e le relative dipendenze.
• Converti un modello Hugging Face nel formato MaxText.
• Esegui un carico di lavoro di addestramento SFT sulla TPU.
• Converti di nuovo il modello ottimizzato nel formato Hugging Face per la pubblicazione.

Costi

In questo documento vengono utilizzati i seguenti componenti fatturabili di Google Cloud:

• Cloud TPU

Per generare una stima dei costi in base all'utilizzo previsto, utilizza il calcolatore prezzi.

I nuovi utenti di Google Cloud potrebbero avere diritto a una prova senza costi.

Al termine delle attività descritte in questo documento, puoi evitare l'addebito di ulteriori costi eliminando le risorse che hai creato. Per saperne di più, consulta Esegui la pulizia.

Importante :per completare questo tutorial, devi avere accesso alle TPU v6e-8. Questo tutorial presuppone che tu abbia già una prenotazione per queste TPU nel tuo progettoGoogle Cloud per assicurarti di avere la capacità disponibile quando ti serve. Per saperne di più sulle opzioni per prenotare la capacità di Cloud TPU per una data e un'ora future, consulta Informazioni sulle prenotazioni Cloud TPU.

Prima di iniziare

• Per utilizzare questo tutorial, è necessario un token di accesso a Hugging Face. Puoi registrarti per un account senza costi su Hugging Face. Dopo aver creato un account, genera un token di accesso:

  1. Nella pagina Benvenuto in Hugging Face, fai clic sull'avatar del tuo account e seleziona Token di accesso.
  2. Nella pagina Token di accesso, fai clic su Crea nuovo token.
  3. Seleziona il tipo di token Lettura e inserisci un nome per il token.
  4. Viene visualizzato il token di accesso. Salva il token in un luogo sicuro.

• Sul sito web di Hugging Face, accetta il contratto di licenza per il modello che prevedi di addestrare. Questo tutorial utilizza il modello gemma3-4b.

Per ottenere le autorizzazioni necessarie per completare questo tutorial, chiedi all'amministratore di concederti i seguenti ruoli IAM nel progetto:

• Amministratore TPU (roles/tpu.admin)
• Utente Service Account (roles/iam.serviceAccountUser)
• Editor Compute (roles/compute.editor)

Per saperne di più sulla concessione dei ruoli, consulta Gestisci l'accesso a progetti, cartelle e organizzazioni.

Potresti anche riuscire a ottenere le autorizzazioni richieste tramite i ruoli personalizzati o altri ruoli predefiniti.

Configura l'ambiente

Imposta le variabili di ambiente eseguendo il seguente script:

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

Sostituisci quanto segue:

• YOUR_PROJECT_ID: il tuo Google Cloud ID progetto
• YOUR_ZONE: la zona che vuoi utilizzare
• YOUR_RESERVATION_NAME: la prenotazione di capacità
• YOUR_TPU_NAME: il nome dell'istanza VM Cloud TPU

Autenticati con Google Cloud eseguendo questo comando:

gcloud auth login

Crea la VM Cloud TPU

Crea un'istanza VM Cloud TPU con 8 chip TPU v6e, associata alla prenotazione della 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}"

Dopo aver creato l'istanza VM, connettiti utilizzando SSH.

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

Completa i seguenti passaggi all'interno dell'istanza VM TPU.

Installare MaxText

Aggiorna i pacchetti di sistema all'interno dell'istanza VM TPU.

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

Installa Python 3.12, richiesto da MaxText, e il relativo pacchetto dell'ambiente virtuale.

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

Utilizza uv per velocizzare l'installazione del pacchetto Python.

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

Crea un ambiente virtuale denominato maxtext_venv e attivalo.

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

Installa MaxText e le dipendenze richieste per le attività post-addestramento.

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

Installa le dipendenze obbligatorie rimanenti eseguendo il seguente comando:

install_tpu_post_train_extra_deps

Converti il modello nel formato MaxText

Per addestrare il modello in formato MaxText, devi convertirlo dal formato Hugging Face al formato MaxText.

Specifica le variabili di ambiente, ad esempio il token di accesso a Hugging Face, il nome del modello che vuoi utilizzare e la directory in cui vuoi salvare il modello in formato 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.

Sostituisci YOUR_HF_TOKEN con il token di accesso a Hugging Face che hai creato in precedenza.

Per convertire il modello dal formato Hugging Face al formato MaxText, esegui lo script seguente. La conversione richiede circa cinque minuti.

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

Avvia il carico di lavoro di addestramento

Una volta completato il processo di conversione, puoi avviare il carico di lavoro SFT.

1. Configura i parametri di addestramento del workload 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. Avvia il job di addestramento. L'operazione richiede circa 10 minuti su un'istanza 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

Converti il modello addestrato di nuovo nel formato Hugging Face

Al termine del carico di lavoro di addestramento, converti nuovamente il modello nel formato Hugging Face.

1. Imposta i percorsi per l'esportazione e i parametri addestrati.

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

2. Esegui la conversione nel formato 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

Una volta completata la conversione, il modello ottimizzato archiviato in /dev/shm/gemma3-4b/hf-trained è pronto per essere utilizzato. Poiché perdi l'accesso ai contenuti della cartella /dev/shm quando la VM viene riavviata, devi spostare il modello ottimizzato in uno spazio di archiviazione permanente o caricarlo su Hugging Face Hub.

Esegui la pulizia

Per evitare addebiti aggiuntivi, elimina le risorse create durante questo tutorial.

Elimina l'istanza VM TPU

Esci dall'istanza VM di Cloud TPU, quindi eliminala.

# 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

Passaggi successivi

• Per saperne di più su Cloud TPU, consulta Introduzione a Cloud TPU.
• Per i dettagli sull'architettura e sulla configurazione della TPU v6e-8, vedi TPU v6e.