Questa pagina è stata tradotta dall'API Cloud Translation.

Addestramento di Llama 3 con PyTorch su TPU v5e

Questo tutorial descrive come addestrare un modello Llama-3-8B utilizzando PyTorch/XLA su TPU v5e utilizzando il set di dati WikiText. Per i dettagli del modello, consulta Meta-Llama-3-8B.

Il modello Llama-3-8B è ospitato sulla piattaforma Hugging Face.

Esistono due versioni di Meta-Llama-3-8B, una da utilizzare con Transformers e un'altra con il codice sorgente originale di Llama 3. Questo tutorial utilizza la versione Transformers perché:

Si integra perfettamente con l'ecosistema Hugging Face: in questo modo è più facile perfezionare il modello, utilizzare pipeline predefinite e accedere a una vasta raccolta di set di dati e strumenti.
Offre flessibilità e personalizzazione: la versione Transformers offre opzioni di personalizzazione e flessibilità significative per la messa a punto e il deployment del modello.
Fornisce assistenza della community: la community di Hugging Face offre documentazione, tutorial e assistenza approfonditi per l'utilizzo dei modelli Transformers.

Per ulteriori informazioni sui trasformatori, consulta la documentazione di Hugging Face Transformers.

Per accedere e utilizzare il modello Meta-Llama-3-8B, incluso il download dei relativi parametri e del tokenizzatore, è necessario un token di accesso utente Hugging Face. Il token fornisce:

Autenticazione e autorizzazione: il token di accesso funge da credenziale e consente ai server di Hugging Face di autorizzare il tuo accesso alle risorse del modello. In questo modo, solo gli utenti autorizzati possono scaricare e utilizzare il modello.
Sicurezza: Hugging Face utilizza i token di accesso per proteggere i propri modelli e impedire accessi o usi impropri non autorizzati.

Per informazioni sulla creazione e sull'utilizzo di un token di accesso per questo tutorial, consulta Eseguire il modello. Per informazioni più complete sulla creazione e sull'utilizzo dei token di accesso, consulta la documentazione di Hugging Face sui token di accesso utente.

Devi anche disporre dell'autorizzazione per accedere al modello Llama 3 8B su Hugging Face. Per ottenere questa autorizzazione, vai al modello Meta-Llama-3-8B su HuggingFace e richiedi l'accesso.

Preparativi per il provisioning di un pod TPU v5lite-16

Questo tutorial è stato testato utilizzando le seguenti variabili di ambiente Cloud TPU. Puoi utilizzare altre variabili per eseguire il provisioning della TPU, purché il tipo di acceleratore, la zona e la versione del runtime siano compatibili. Ad esempio, in questo tutorial, europe-west4-b viene utilizzato come zona. Puoi utilizzare qualsiasi altra zona che supporti la versione TPU (tipo di acceleratore) in esecuzione (v5litepod-16 in questo tutorial).

Imposta le seguenti variabili di ambiente della VM TPU.

   export TPU_NAME=queued-resources-node-id #The TPU name is the queued resource node-id
   export PROJECT_ID=your-project-id
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=europe-west4-b
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export QUEUED_RESOURCE_ID=queued-resource-id
   export VALID_UNTIL_DURATION=1d

Quando hai accesso al modello Meta-Llama-3-8B su Hugging Face, prepara l'ambiente TPU per eseguire il tutorial.

Segui la guida sulla configurazione dell'ambiente Cloud TPU per assicurarti di disporre dell'accesso appropriato per utilizzare le Cloud TPU.

Crea un'identità di servizio per la VM TPU.

gcloud alpha compute tpus tpu-vm service-identity create --zone=zone

Crea un account di servizio TPU e concedi accesso ai servizi Google Cloud .

Gli account di servizio consentono al Google Cloud servizio TPU di accedere ad altri Google Cloud servizi. È consigliato un account di servizio gestito dall'utente. Puoi creare un account di servizio dalla console Google Cloud o tramite il comandogcloud.

Crea un account di servizio utilizzando lo strumento a riga di comando gcloud:
```
gcloud iam service-accounts create your-service-account-name \
--description="your-sa-description" \
--display-name="your-sa-display-name"
export SERVICE_ACCOUNT_NAME=your-service-account-name
```
Crea un account di servizio dalla console Google Cloud:
1. Vai alla pagina Account di servizio nella console Google Cloud.
2. Fai clic su Crea account di servizio.
3. Inserisci il nome dell'account di servizio.
4. (Facoltativo) Inserisci una descrizione per l'account di servizio.
5. Fai clic su Crea e continua.
6. Scegli i ruoli da concedere all'account di servizio.
7. Fai clic su Continua.
8. (Facoltativo) Specifica gli utenti o i gruppi che possono gestire l'account di servizio.
9. Fai clic su Fine per completare la creazione del account di servizio.
Dopo aver creato l'account di servizio, segui questi passaggi per concedere i ruoli dell'account di servizio.

Sono necessari i seguenti ruoli:
- TPU Admin: necessario per creare una TPU
- Amministratore Storage: necessario per accedere a Cloud Storage
- Writer log
- Monitoring Metric Writer: necessario per scrivere le metriche in Cloud Monitoring
L'amministratore deve concederti il ruolo roles/resourcemanager.projectIamAdmin per consentirti di assegnare i ruoli IAM agli utenti. Anche un utente con il ruolo Amministratore IAM progetto roles/resourcemanager.projectIamAdmin può assegnare questo ruolo.

Utilizza i seguenti comandi gcloud per aggiungere i ruoli dell'account di servizio:
```
gcloud projects add-iam-policy-binding ${PROJECT_ID} \
   --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
   --role roles/tpu.admin
gcloud projects add-iam-policy-binding ${PROJECT_ID} \
   --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
   --role roles/storage.admin
gcloud projects add-iam-policy-binding ${PROJECT_ID} \
   --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
   --role roles/logging.logWriter
gcloud projects add-iam-policy-binding ${PROJECT_ID} \
   --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
   --role roles/monitoring.metricWriter
```
Puoi assegnare i ruoli anche utilizzando la console Google Cloud.

Nella console Google Cloud, seleziona i seguenti ruoli:
1. Seleziona il tuo account di servizio e fai clic su Aggiungi principale.
2. Nel campo Nuove entità, inserisci l'indirizzo email del tuo account di servizio.
3. Nel menu a discesa Seleziona un ruolo, cerca il ruolo (ad esempio Amministratore Storage) e selezionalo.
4. Fai clic su Salva.
Esegui l'autenticazione con Google Cloud e configura il progetto e la zona predefiniti per Google Cloud CLI.
```
gcloud auth login
gcloud config set project PROJECT_ID
gcloud config set compute/zone ZONE
```

Capacità sicura

Quando è tutto pronto per assicurarti la capacità TPU, consulta la pagina delle quote per saperne di più sul sistema delle quote Cloud. Se hai altre domande su come assicurarti la capacità, contatta il team di vendita o dell'account Cloud TPU.

Esegui il provisioning dell'ambiente Cloud TPU

Puoi eseguire il provisioning delle VM TPU con GKE, con GKE e XPK, o come risorse in coda.

Prerequisiti

Questo tutorial è stato testato con Python 3.10 o versioni successive.
Verifica che il tuo progetto disponga di una quota TPUS_PER_TPU_FAMILY sufficiente, che specifica il numero massimo di chip a cui puoi accedere all'interno del progettoGoogle Cloud .
Verifica che il tuo progetto disponga di una quota TPU sufficiente per:
- Quota VM TPU
- Quota di indirizzi IP
- Quota Hyperdisk bilanciata
Autorizzazioni del progetto per gli utenti
- Se utilizzi GKE con XPK, consulta Autorizzazioni di Cloud Console per l'account utente o di servizio per conoscere le autorizzazioni necessarie per eseguire XPK.

Provisiona un TPU v5litepod-16

Crea una VM TPU:

gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID}   \
--node-id=${TPU_NAME}  \
--project=${PROJECT_ID}   \
--zone=${ZONE}   \
--accelerator-type=${ACCELERATOR_TYPE}   \
--runtime-version=${RUNTIME_VERSION}   \
--service-account=${SERVICE_ACCOUNT_NAME}   \
--spot

Verifica che la TPU sia nello stato ACTIVE:

gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
--project=${PROJECT_ID} \
--zone=${ZONE}

Quando la TPU diventa attiva (ACTIVE), viene visualizzato un output simile al seguente:

createTime: '2025-02-28T21:16:08.053492925Z'
name: projects/my-project/locations/zone/queuedResources/tpu-name-zone
spot: {}
state:
  state: ACTIVE
tpu:
  nodeSpec:
  - node:
      acceleratorType: v5litepod-16
      networkConfig:
        enableExternalIps: true
        network: default
      queuedResource: projects/19672137403/locations/zone/queuedResources/qr-name
      runtimeVersion: v2-alpha-tpuv5-lite
      schedulingConfig: {}
      my-service-account@your-project-id.iam.gserviceaccount.com
        email: 19672137854-compute@developer.iam.gserviceaccount.com
      shieldedInstanceConfig: {}
    nodeId: tpu-name
    parent: projects/19672137403/locations/zone

Installazione

Installa il pytorch-tpu/transformers fork di Hugging Face Transformers e le dipendenze. Questo tutorial è stato testato con le seguenti versioni delle dipendenze:

torch: compatibile con 2.6.0
torch_xla[tpu]: compatibile con 2.6.0
jax: 0.4.38
jaxlib: 0.4.38

Installa il software e le dipendenze del framework

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} \
    --zone=${ZONE} \
    --worker=all \
    --command='git clone -b flash_attention https://github.com/pytorch-tpu/transformers.git
    sudo apt install python3.10-venv
    python -m venv /home/$USER/venv/
    source ~/venv/bin/activate
    cd transformers
    pip3 install --user -e .
    pip3 install datasets
    pip3 install evaluate
    pip3 install scikit-learn
    pip3 install accelerate
    pip install torch~=2.6.0 torch_xla[tpu]~=2.6.0 -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html
    pip install jax==0.4.38 jaxlib==0.4.38 -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html'

Al termine dell'installazione, visualizzerai un output simile al seguente:

Collecting jax==0.4.38
  Downloading jax-0.4.38-py3-none-any.whl (2.1 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.1/2.1 MB 18.0 MB/s eta 0:00:00
Collecting jaxlib==0.4.38
  Downloading jaxlib-0.4.38-cp310-cp310-manylinux2014_x86_64.whl (85.0 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 85.0/85.0 MB 10.1 MB/s eta 0:00:00
Successfully installed jax-0.4.38 jaxlib-0.4.38 ml-dtypes-0.5.1 opt-einsum-3.4.0
Collecting opt-einsum
  Downloading opt_einsum-3.4.0-py3-none-any.whl (71 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 71.9/71.9 KB 186.4 kB/s eta 0:00:00
Requirement already satisfied: numpy>=1.24 in /home/your-username/.local/lib/python3.10/site-packages (from jax==0.4.38) (2.2.3)
Requirement already satisfied: scipy>=1.10 in /home/your-username/.local/lib/python3.10/site-packages (from jax==0.4.38) (1.15.2)
Successfully installed jax-0.4.38 jaxlib-0.4.38 ml-dtypes-0.5.1 opt-einsum-3.4.0
Successfully installed jax-0.4.38 jaxlib-0.4.38 ml-dtypes-0.5.1 opt-einsum-3.4.0
Collecting ml-dtypes>=0.2.0
  Downloading ml_dtypes-0.5.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (4.7 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 4.7/4.7 MB 13.8 MB/s eta 0:00:00
Installing collected packages: opt-einsum, ml-dtypes, jaxlib, jax
Successfully installed jax-0.4.38 jaxlib-0.4.38 ml-dtypes-0.5.1 opt-einsum-3.4.0

Configurare le configurazioni del modello

Il comando di addestramento nella sezione successiva, Esegui il modello, utilizza due file di configurazione JSON per definire i parametri del modello e la configurazione FSDP (Fully Sharded Data Parallel). La suddivisione FSDP viene utilizzata per adattare i pesi del modello a un batch di dimensioni maggiori durante l'addestramento. Quando si esegue l'addestramento con modelli più piccoli, potrebbe essere sufficiente utilizzare il parallelismo dei dati e replicare i pesi su ogni dispositivo. Per ulteriori informazioni su come suddividere i tensori su più dispositivi in PyTorch/XLA, consulta la Guida dell'utente SPMD di PyTorch/XLA.

Questo comando crea il file di configurazione dei parametri del modello per Llama3-8B. Per altri modelli, trova la configurazione su Hugging Face. Ad esempio, consulta Llama2-7B config.

cat > llama-config.json <<EOF
{
    "architectures": [
        "LlamaForCausalLM"
    ],
    "attention_bias": false,
    "attention_dropout": 0.0,
    "bos_token_id": 128000,
    "eos_token_id": 128001,
    "hidden_act": "silu",
    "hidden_size": 4096,
    "initializer_range": 0.02,
    "intermediate_size": 14336,
    "max_position_embeddings": 8192,
    "model_type": "llama",
    "num_attention_heads": 32,
    "num_hidden_layers": 32,
    "num_key_value_heads": 8,
    "pretraining_tp": 1,
    "rms_norm_eps": 1e-05,
    "rope_scaling": null,
    "rope_theta": 500000.0,
    "tie_word_embeddings": false,
    "torch_dtype": "bfloat16",
    "transformers_version": "4.40.0.dev0",
    "use_cache": false,
    "vocab_size": 128256
}
EOF

Crea il file di configurazione FSDP:

cat > fsdp-config.json <<EOF
{
    "fsdp_transformer_layer_cls_to_wrap": [
        "LlamaDecoderLayer"
    ],
    "xla": true,
    "xla_fsdp_v2": true,
    "xla_fsdp_grad_ckpt": true
}
EOF

Per ulteriori informazioni su FSDP, consulta FSDPv2.

Carica i file di configurazione nelle VM TPU utilizzando i seguenti comandi:

 ssh-add ~/.ssh/google_compute_engine #Setup SSH Key in the SSH agent.
 gcloud alpha compute tpus tpu-vm scp llama-config.json fsdp-config.json $TPU_NAME:. \
    --worker=all \
    --project=${PROJECT_ID} \
    --zone=${ZONE}

Questo comando genererà un output simile al seguente:

Using scp batch size of 4.Attempting to SCP into 1 nodes with a total of 4 workers.
SCP: Attempting to connect to worker 0...
SCP: Attempting to connect to worker 1...
SCP: Attempting to connect to worker 2...
SCP: Attempting to connect to worker 3...
llama-config.json                    100%  707     4.1KB/s   00:00
llama-config.json                    100%  707     4.0KB/s   00:00
llama-config.json                    100%  707     4.1KB/s   00:00
llama-config.json                    100%  707     4.1KB/s   00:00
fsdp-config.json                     100%  156     0.9KB/s   00:00
fsdp-config.json                     100%  156     0.9KB/s   00:00
fsdp-config.json                     100%  156     0.9KB/s   00:00
fsdp-config.json                     100%  156     0.9KB/s   00:00

Esegui il modello

Utilizzando i file di configurazione creati nella sezione precedente, esegui lo script run_clm.py per addestrare il modello Llama 3 8B sul set di dati WikiText. L'esecuzione dello script di addestramento su una TPU v5litepod-16 richiede circa 10 minuti.

Genera un nuovo token Hugging Face se non ne hai già uno:
1. Fai clic su Il tuo profilo > Impostazioni > Token di accesso.
2. Seleziona Nuovo token.
3. Specifica un nome a tua scelta e un ruolo di almeno Lettura.
4. Seleziona Genera un token.
Utilizza il token Hugging Face per accedere a Hugging Face dalla VM TPU utilizzando il seguente comando.

Sostituisci la variabile token huggingface-cli login con quella generata da Hugging Face nel passaggio precedente:
```
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='
pip install -U "huggingface_hub[cli]"
export PATH="/home/$USER/.local/bin/:$PATH"
huggingface-cli login --token hf_abcxyzEFg'
```
Questo comando ti consente di accedere a Hugging Face e di visualizzare il token attivo corrente.

Esegui l'addestramento del modello:

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} \
    --zone=${ZONE} \
    --worker=all \
    --command='
    source ~/venv/bin/activate
    export PJRT_DEVICE=TPU
    export XLA_USE_SPMD=1
    export ENABLE_PJRT_COMPATIBILITY=true
        # Optional variables for debugging:
    export XLA_IR_DEBUG=1
    export XLA_HLO_DEBUG=1
    export PROFILE_EPOCH=0
    export PROFILE_STEP=3
    export PROFILE_DURATION_MS=100000
        # Set PROFILE_LOGDIR to a local VM path or gs://my-bucket/profile_path
    export PROFILE_LOGDIR=your-bucket/profile_path
    python3 transformers/examples/pytorch/language-modeling/run_clm.py \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --per_device_train_batch_size 16 \
    --do_train \
    --output_dir /home/$USER/tmp/test-clm \
    --overwrite_output_dir \
    --config_name /home/$USER/llama-config.json \
    --cache_dir /home/$USER/cache \
    --tokenizer_name meta-llama/Meta-Llama-3-8B \
    --block_size 8192 \
    --optim adafactor \
    --save_strategy no \
    --logging_strategy no \
    --fsdp "full_shard" \
    --fsdp_config /home/$USER/fsdp-config.json \
    --torch_dtype bfloat16 \
    --dataloader_drop_last yes \
    --flash_attention \
    --max_steps 20'

Il passaggio di addestramento richiede circa 10 minuti. Verso la fine dell'addestramento, vedrai messaggi simili ai seguenti:

[INFO|trainer.py:2053] 2025-03-18 22:05:02,536 >> ***** Running training *****
[INFO|trainer.py:2054] 2025-03-18 22:05:02,536 >>   Num examples = 272
[INFO|trainer.py:2055] 2025-03-18 22:05:02,536 >>   Num Epochs = 2
[INFO|trainer.py:2056] 2025-03-18 22:05:02,536 >>   Instantaneous batch size per device = 16
[INFO|trainer.py:2059] 2025-03-18 22:05:02,536 >>   Total train batch size (w. parallel, distributed & accumulation) = 16
[INFO|trainer.py:2060] 2025-03-18 22:05:02,536 >>   Gradient Accumulation steps = 1
[INFO|trainer.py:2061] 2025-03-18 22:05:02,536 >>   Total optimization steps = 20
[INFO|trainer.py:2062] 2025-03-18 22:05:02,537 >>   Number of trainable parameters = 8,030,261,248
  0%|          | 0/20 [00:00<?, ?it/s][INFO|trainer.py:2143] 2025-03-18 22:05:02,540 >> Profiling server started: <_XLAC.profiler.ProfilerServer object at 0x7f01bdcb6770>
  5%|▌         | 1/20 [00:07<02:29,  7.86s/it]/home/your-username/venv/lib/python3.10/site-packages/torch/nn/modules/module.py:1810: UserWarning: For backward hooks to be called, module output should be a Tensor or a tuple of Tensors but received <class 'transformers.modeling_outputs.CausalLMOutputWithPast'>
  5%|▌         | 1/20 [00:07<02:29,  7.89s/it]Compilation at Step 0, time: 213.83555555343628
/home/your-username/venv/lib/python3.10/site-packages/torch/nn/modules/module.py:1810:
 10%|█         | 2/20 [03:43<38:57, 129.87s/it]Compilation at Step 0, time: 213.12156581878662
/home/your-username/venv/lib/python3.10/site-packages/torch/nn/modules/module.py:"
 10%|█         | 2/20 [03:40<38:29, 128.30s/it]Compilation at Step 1, time: 224.5414960384369
 15%|█▌        | 3/20 [07:22<48:31, 171.24s/it]Compilation at Step 1, time: 226.23664164543152
 15%|█▌        | 3/20 [07:26<48:56, 172.73s/it]Compilation at Step 1, time: 226.9180543422699
Compilation at Step 1, time: 224.3874273300171
 20%|██        | 4/20 [07:23<27:45, 104.10s/it]Starting to trace for 100000 ms. Remaining attempt(s): 2
2025-03-18 22:12:32.104419: W external/tsl/tsl/profiler/lib/profiler_session.cc:109] Profiling is late by 847930 nanoseconds and will start immediately.
Starting to trace for 100000 ms. Remaining attempt(s): 2
2025-03-18 22:12:32.104373: W external/tsl/tsl/profiler/lib/profiler_session.cc:109] Profiling is late by 763960 nanoseconds and will start immediately.
Starting to trace for 100000 ms. Remaining attempt(s): 2
2025-03-18 22:12:32.104538: W external/tsl/tsl/profiler/lib/profiler_session.cc:109] Profiling is late by 854020 nanoseconds and will start immediately.
2025-03-18 22:12:32.104347: W external/tsl/tsl/profiler/lib/profiler_session.cc:109] Profiling is late by 761070 nanoseconds and will start immediately.
Starting to trace for 100000 ms. Remaining attempt(s): 2
 85%|████████▌ | 17/20 [07:55<00:06,  2.26s/it]Compilation at Step -1, time: 3.676558494567871
Compilation at Step -1, time: 3.447533130645752
Compilation at Step -1, time: 3.5890843868255615
Compilation at Step -1, time: 3.4956483840942383
100%|██████████| 20/20 [11:39<00:00, 35.14s/it][INFO|trainer.py:2350] 2025-03-18 22:16:42,476 >>
Training completed. Do not forget to share your model on huggingface.co/models =)

100%|██████████| 20/20 [11:47<00:00, 35.23s/it][INFO|trainer.py:2350] 2025-03-18 22:16:43,239 >>

Training completed. Do not forget to share your model on huggingface.co/models =)

Esegui la pulizia

Al termine dell'addestramento, segui questo passaggio per eliminare la risorsa in coda e la VM TPU. In questo modo, la fatturazione per l'utilizzo della VM TPU verrà interrotta.

  gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
       --project=${PROJECT_ID} \
       --zone=${ZONE} \
       --force \
       --async