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

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

Riepilogo del carico di lavoro

L'esecuzione di questo carico di lavoro SFT su una VM Cloud TPU comporta i seguenti componenti:

• Modello: Gemma 3 4B (gemma3-4b), un modello linguistico leggero con pesi aperti. In questo tutorial, i pesi del modello vengono convertiti dal formato Hugging Face al formato compatibile con MaxText (mt-format) e preparati per l'esecuzione su una singola VM Cloud TPU v6e-8 utilizzando JAX.
• Metodo di addestramento: SFT utilizzando MaxText. Il carico di lavoro esegue una pipeline di ottimizzazione su una singola istanza Cloud TPU v6e-8 utilizzando JAX. Ottimizza i parametri del modello in oltre 1000 passaggi utilizzando l'addestratore SFT integrato di MaxText (maxtext.trainers.post_train.sft.train_sft) e i set di dati di ottimizzazione delle istruzioni incentrati sul dialogo.
• Set di dati e modello: UltraChat 200k (HuggingFaceH4/ultrachat_200k), un set di dati di dialogo multi-turn di alta qualità su larga scala (train_sft suddiviso con la colonna messages). Viene utilizzato per ottimizzare il modello e migliorare le sue capacità di conversazione e di seguire le istruzioni.
• Infrastruttura e framework: una singola istanza VM Cloud TPU v6e-8 (ct6e-standard-8t). Il framework è MaxText (maxtext[tpu-post-train]) basato su JAX.

Obiettivi

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

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 Google Cloud utenti 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 Google Cloud progetto per assicurarti di avere la capacità disponibile quando ne hai bisogno. Per saperne di più sulle opzioni per prenotare la capacità Cloud TPU per una data e un'ora future, consulta Informazioni sulle prenotazioni Cloud TPU.

Prima di iniziare

• Per utilizzare questo tutorial, devi disporre di 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 intendi 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 tuo progetto:

• Amministratore TPU (roles/tpu.admin)
• Utente dell'account di servizio (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: l' Google Cloud ID progetto
• YOUR_ZONE: la zona che vuoi utilizzare
• YOUR_RESERVATION_NAME: la prenotazione della capacità
• YOUR_TPU_NAME: il nome dell'istanza VM Cloud TPU

Esegui l'autenticazione con Google Cloud eseguendo il seguente comando:

gcloud auth login

Crea la VM Cloud TPU

Crea un'istanza VM Cloud TPU con 8 chip TPU v6e, associati 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 passaggi seguenti all'interno dell'istanza VM TPU.

Installa 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 di 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 rimanenti richieste eseguendo il seguente comando:

install_tpu_post_train_extra_deps

Converti il modello in formato MaxText

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

Specifica le variabili di ambiente, come 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 il seguente script. 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

Al termine del processo di conversione, puoi avviare il carico di lavoro SFT.

1. Configura i parametri di addestramento del carico di lavoro 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. Questa 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 di nuovo il modello addestrato in formato Hugging Face

Al termine del carico di lavoro di addestramento, converti di nuovo il modello in 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 di nuovo in 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

Al termine della conversione, il modello ottimizzato memorizzato in /dev/shm/gemma3-4b/hf-trained è pronto per essere utilizzato. Poiché perdi l'accesso ai contenuti della cartella /dev/shm al riavvio della VM, 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 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 v6e-8 TPU, consulta TPU v6e.