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 MaxText, uno stack di addestramento basato su JAX ad alte prestazioni per modelli linguistici di grandi dimensioni (LLM). Google Cloud

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:

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 Liberare spazio.

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. Una volta 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:

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="ZONE_NAME"
export RESERVATION="RESERVATION_NAME"
export NAME="TPU_MACHINE_NAME"

Sostituisci quanto segue:

  • YOUR_PROJECT_ID: il tuo Google Cloud ID progetto
  • ZONE_NAME: la zona che vuoi utilizzare
  • RESERVATION_NAME: la tua prenotazione della capacità
  • TPU_MACHINE_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, associata alla prenotazione della capacità.

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

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

gcloud compute tpus tpu-vm 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 python3.12 python3.12-venv

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 venv --python 3.12 --seed maxtext_venv
source maxtext_venv/bin/activate

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

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

Installa le dipendenze richieste rimanenti eseguendo il seguente comando:

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

Libera spazio

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

Elimina l'istanza VM TPU

Esci dall'istanza VM Cloud TPU ed eliminala.

gcloud alpha compute tpus tpu-vm delete $NAME --zone=$ZONE --project=$PROJECT --quiet

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.