Pubblica modelli aperti Llama 3 utilizzando Cloud TPU multi-host su Vertex AI Prediction con Saxml

Llama 3 è un modello linguistico di grandi dimensioni (LLM) open source di Meta. Questa guida mostra come pubblicare un LLM Llama 3 utilizzando Tensor Processing Unit (TPU) multi-host su Vertex AI Prediction con Saxml.

In questa guida, scarichi i pesi e il tokenizzatore del modello Llama 3 70B e li esegui su Vertex AI Prediction, che esegue Saxml su TPU.

Prima di iniziare

Ti consigliamo di utilizzare una VM M2 ottimizzato per la memoria per scaricare il modello e convertirlo in Saxml. Questo perché il processo di conversione del modello richiede una quantità significativa di memoria e potrebbe non riuscire se scegli un tipo di macchina con memoria insufficiente.

  1. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.
  2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  3. Make sure that billing is enabled for your Google Cloud project.

  4. Enable the Vertex AI and Artifact Registry APIs.

    Enable the APIs

  5. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  6. Make sure that billing is enabled for your Google Cloud project.

  7. Enable the Vertex AI and Artifact Registry APIs.

    Enable the APIs

  8. In the Google Cloud console, activate Cloud Shell.

    Activate Cloud Shell

    At the bottom of the Google Cloud console, a Cloud Shell session starts and displays a command-line prompt. Cloud Shell is a shell environment with the Google Cloud CLI already installed and with values already set for your current project. It can take a few seconds for the session to initialize.

  9. Segui la documentazione di Artifact Registry per installare Docker.
  10. Assicurati di disporre di quote sufficienti per 16 chip TPU v5e per Vertex AI Prediction.

Questo tutorial presuppone che tu stia utilizzando Cloud Shell per interagire con Google Cloud. Se vuoi utilizzare una shell diversa da Cloud Shell, esegui la seguente configurazione aggiuntiva:

  1. Install the Google Cloud CLI.
  2. To initialize the gcloud CLI, run the following command:

    gcloud init

Se utilizzi una shell diversa da Cloud Shell per il deployment del modello, assicurati che la versione di Google Cloud CLI sia successiva alla 475.0.0. Puoi aggiornare Google Cloud CLI eseguendo il comando gcloud components update.

Se esegui il deployment del modello utilizzando l'SDK Vertex AI, assicurati di avere la versione 1.50.0 o successiva.

Accedi al modello e scarica i relativi pesi

I passaggi riportati di seguito si riferiscono a un'istanza di Vertex AI Workbench con una VM M2 ottimizzato per la memoria. Per informazioni su come modificare il tipo di macchina di un'istanza di Vertex AI Workbench, consulta Modificare il tipo di macchina di un'istanza di Vertex AI Workbench.

  1. Vai alla pagina del consenso per il modello Llama.

  2. Seleziona Llama 3, compila il modulo per il consenso e accetta i Termini e condizioni.

  3. Controlla se nella posta in arrivo è presente un'email contenente un URL firmato.

  4. Scarica lo script download.sh da GitHub eseguendo il seguente comando:

    wget https://raw.githubusercontent.com/meta-llama/llama3/main/download.sh
    chmod +x download.sh
    
  5. Per scaricare i pesi del modello, esegui lo script download.sh che hai scaricato da GitHub.

  6. Quando richiesto, inserisci l'URL firmato dell'email che hai ricevuto nella sezione precedente.

  7. Quando ti viene chiesto di scaricare i modelli, inserisci 70B.

Converti i pesi del modello in formato Saxml

  1. Esegui il seguente comando per scaricare Saxml:

    git clone https://github.com/google/saxml.git
    
  2. Esegui i comandi seguenti per configurare un ambiente virtuale Python:

    python -m venv .
    source bin/activate
    
  3. Esegui i seguenti comandi per installare le dipendenze:

    pip install --upgrade pip
    
    pip install paxml
    
    pip install praxis
    
    pip install torch
    
  4. Per convertire i pesi del modello in formato Saxml, esegui il seguente comando:

    python3 saxml/saxml/tools/convert_llama_ckpt.py \
        --base PATH_TO_META_LLAMA3 \
        --pax PATH_TO_PAX_LLAMA3 \
        --model-size llama3_70b
    

    Sostituisci quanto segue:

    • PATH_TO_META_LLAMA3: il percorso della directory contenente le ponderazioni del modello scaricate
    • PATH_TO_PAX_LLAMA3: il percorso della directory in cui memorizzare i pesi del modello convertito

    I modelli convertiti verranno inseriti nella $PATH_TO_PAX_LLAMA3/checkpoint_00000000.

  5. Copia il file del tokenizzatore dalla directory originale in una sottocartella denominata vocabs come segue:

    cp $PATH_TO_META_LLAMA3/tokenizer.model $PATH_TO_PAX_LLAMA3/vocabs/tokenizer.model
    
  6. Aggiungi un file commit_success.txt vuoto nella cartella $PATH_TO_PAX_LLAMA3 e le sottocartelle metadata e state al suo interno come segue:

    touch $PATH_TO_PAX_LLAMA3/checkpoint_00000000/commit_success.txt
    touch $PATH_TO_PAX_LLAMA3/checkpoint_00000000/metadata/commit_success.txt
    touch $PATH_TO_PAX_LLAMA3/checkpoint_00000000/state/commit_success.txt
    

    La cartella $PATH_TO_PAX_LLAMA3 ora contiene le seguenti cartelle e file:

    $PATH_TO_PAX_LLAMA3/checkpoint_00000000/commit_success.txt
    $PATH_TO_PAX_LLAMA3/checkpoint_00000000/metadata/
    $PATH_TO_PAX_LLAMA3/checkpoint_00000000/metadata/commit_success.txt
    $PATH_TO_PAX_LLAMA3/checkpoint_00000000/state/
    $PATH_TO_PAX_LLAMA3/checkpoint_00000000/state/commit_success.txt
    $PATH_TO_PAX_LLAMA3/vocabs/tokenizer.model
    

Crea un bucket Cloud Storage

Crea un bucket Cloud Storage per archiviare i pesi del modello convertito.

  1. In Cloud Shell, esegui i seguenti comandi, sostituendo PROJECT_ID con l'ID progetto:

    projectid=PROJECT_ID
    gcloud config set project ${projectid}
    
  2. Per creare il bucket, esegui il seguente comando:

    gcloud storage buckets create gs://WEIGHTS_BUCKET_NAME
    

    Sostituisci WEIGHTS_BUCKET_NAME con il nome che vuoi utilizzare per il bucket.

Copia i pesi del modello nel bucket Cloud Storage

Per copiare i pesi del modello nel bucket, esegui il seguente comando:

gcloud storage cp PATH_TO_PAX_LLAMA3/* gs://WEIGHTS_BUCKET_NAME/llama3/pax_70b/ --recursive

Carica il modello

Un container Saxml predefinito è disponibile all'indirizzo us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest.

Per caricare una risorsa Model in Vertex AI Prediction utilizzando il contenitore predefinito Saxml, esegui il comando gcloud ai models upload come segue:

gcloud ai models upload \
    --region=LOCATION \
    --display-name=MODEL_DISPLAY_NAME \
    --container-image-uri=us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest \
    --artifact-uri='gs://WEIGHTS_BUCKET_NAME/llama3/pax_70b' \
    --container-args='--model_path=saxml.server.pax.lm.params.lm_cloud.LLaMA3_70BFP16x16' \
    --container-args='--platform_chip=tpuv5e' \
    --container-args='--platform_topology=4x4' \
    --container-args='--ckpt_path_suffix=checkpoint_00000000' \
    --container-deployment-timeout-seconds=2700 \
    --container-ports=8502 \
    --project=PROJECT_ID

Effettua le seguenti sostituzioni:

  • LOCATION: la regione in cui utilizzi Vertex AI. Tieni presente che le TPU sono disponibili solo in us-west1.
  • MODEL_DISPLAY_NAME: il nome visualizzato che vuoi assegnare al modello
  • PROJECT_ID: l'ID del tuo progetto Google Cloud

Creare un endpoint di previsione online

Per creare l'endpoint, esegui il seguente comando:

gcloud ai endpoints create \
    --region=LOCATION \
    --display-name=ENDPOINT_DISPLAY_NAME \
    --project=PROJECT_ID

Sostituisci ENDPOINT_DISPLAY_NAME con il nome visualizzato che vuoi assegnare al tuo endpoint.

Esegui il deployment del modello nell'endpoint

Una volta che l'endpoint è pronto, esegui il deployment del modello.

In questo tutorial esegui il deployment di un modello Llama 3 70B suddiviso in 16 chip Cloud TPU v5e utilizzando una topologia 4x4. Tuttavia, puoi specificare una delle seguenti topologie Cloud TPU multi-host supportate:

Tipo di macchina Topologia Numero di chip TPU Numero di host
ct5lp-hightpu-4t 4x4 16 2
ct5lp-hightpu-4t 4x8 32 4
ct5lp-hightpu-4t 8x8 64 8
ct5lp-hightpu-4t 8x16 128 16
ct5lp-hightpu-4t 16x16 256 32

Se stai implementando un modello Llama diverso definito nel repo GitHub di Saxml, assicurati che sia partizionato in base al numero di dispositivi di destinazione e che Cloud TPU abbia memoria sufficiente per caricare il modello.

Per informazioni sul deployment di un modello su Cloud TPU con un solo host, consulta Eseguire il deployment di un modello.

Per ulteriori informazioni sui tipi di Cloud TPU v5e, consulta TPU v5e.

  1. Ottieni l'ID endpoint per l'endpoint di previsione online:

    ENDPOINT_ID=$(gcloud ai endpoints list \
        --region=LOCATION \
        --filter=display_name=ENDPOINT_NAME \
        --format="value(name)")
    
  2. Recupera l'ID modello del tuo modello:

    MODEL_ID=$(gcloud ai models list \
        --region=LOCATION \
        --filter=display_name=DEPLOYED_MODEL_NAME \
        --format="value(name)")
    
  3. Esegui il deployment del modello nell'endpoint:

    gcloud ai endpoints deploy-model $ENDPOINT_ID \
        --region=LOCATION \
        --model=$MODEL_ID \
        --display-name=DEPLOYED_MODEL_NAME \
        --machine-type=ct5lp-hightpu-4t \
        --tpu-topology=4x4 \
        --traffic-split=0=100
    

    Sostituisci DEPLOYED_MODEL_NAME con un nome per il deployment. Può essere uguale al nome visualizzato del modello (MODEL_DISPLAY_NAME).

    L'operazione di deployment potrebbe scadere.

    Il comando deploy-model restituisce un ID operazione che può essere utilizzato per verificare quando l'operazione è stata completata. Puoi eseguire il polling per lo stato dell'operazione finché la risposta non include "done": true. Utilizza il seguente comando per eseguire il polling dello stato:

    gcloud ai operations describe \
    --region=LOCATION \
    OPERATION_ID
    

    Sostituisci OPERATION_ID con l'ID operazione restituito dal comando precedente.

Ottenere previsioni online dal modello di cui è stato eseguito il deployment

Per ottenere le previsioni online dall'endpoint di previsione Vertex AI, esegui il gcloud ai endpoints predict comando.

  1. Esegui il seguente comando per creare un file request.json contenente una richiesta di previsione di esempio:

    cat << EOF > request.json
    {"instances": [{"text_batch": "the distance between Earth and Moon is "}]}
    EOF
    
  2. Per inviare la richiesta di previsione online all'endpoint, esegui il seguente comando:

    gcloud ai endpoints predict $ENDPOINT_ID \
        --project=PROJECT_ID \
        --region=LOCATION \
        --json-request=request.json
    

Esegui la pulizia

Per evitare ulteriori costi di Vertex AI, elimina le risorse Google Cloud che hai creato durante questo tutorial:

  1. Per annullare il deployment del modello dall'endpoint ed eliminare l'endpoint, esegui i seguenti comandi:

    ENDPOINT_ID=$(gcloud ai endpoints list \
       --region=LOCATION \
       --filter=display_name=ENDPOINT_NAME \
       --format="value(name)")
    
    DEPLOYED_MODEL_ID=$(gcloud ai endpoints describe $ENDPOINT_ID \
       --region=LOCATION \
       --format="value(deployedModels.id)")
    
    gcloud ai endpoints undeploy-model $ENDPOINT_ID \
      --region=LOCATION \
      --deployed-model-id=$DEPLOYED_MODEL_ID
    
    gcloud ai endpoints delete $ENDPOINT_ID \
       --region=LOCATION \
       --quiet
    
  2. Per eliminare il modello, esegui i seguenti comandi:

    MODEL_ID=$(gcloud ai models list \
       --region=LOCATION \
       --filter=display_name=DEPLOYED_MODEL_NAME \
       --format="value(name)")
    
    gcloud ai models delete $MODEL_ID \
       --region=LOCATION \
       --quiet