Questa pagina è stata tradotta dall'API Cloud Translation.

Scala i carichi di lavoro ML utilizzando Ray

Introduzione

Lo strumento Ray per Cloud TPU combina l'API Cloud TPU e Ray Jobs con l'obiettivo di migliorare l'esperienza di sviluppo degli utenti su Cloud TPU. Questa guida utente fornisce un esempio minimo di come utilizzare Ray con le Cloud TPU. Questi esempi non sono pensati per essere utilizzati in produzione servizi e sono a scopo puramente illustrativo.

Cosa include questo strumento?

Per comodità, lo strumento fornisce:

Astrazioni generiche che nascondono il boilerplate per le azioni TPU comuni
Esempi di giocattoli che puoi creare per i tuoi flussi di lavoro di base

In particolare:

tpu_api.py: Wrapper Python per le operazioni TPU di base utilizzando l'API Cloud TPU.
tpu_controller.py: representatione di classe di una TPU. Si tratta essenzialmente di un wrapper per tpu_api.py.
ray_tpu_controller.py: Controller TPU con funzionalità Ray. che elimina il boilerplate per Cluster Ray e job Ray.
run_basic_jax.py: esempio di base che mostra come utilizzare RayTpuController per print(jax.device_count()).
run_hp_search.py: esempio di base che mostra come utilizzare Ray Tune con JAX/Flax su MNIST.
run_pax_autoresume.py: Esempio che mostra come utilizzare RayTpuController per errore con PAX come carico di lavoro di esempio.

Configurazione del nodo principale del cluster Ray

Uno dei modi di base in cui puoi utilizzare Ray con un pod di TPU è configurare il pod di TPU come ammasso Ray. La creazione di una VM CPU separata come VM coordinatore è naturale per farlo. L'immagine seguente mostra un esempio di ammasso Ray configurazione:

Esempio di configurazione di un cluster Ray

I seguenti comandi mostrano come configurare un cluster Ray utilizzando Google Cloud CLI:

$ gcloud compute instances create my_tpu_admin --machine-type=n1-standard-4 ...
$ gcloud compute ssh my_tpu_admin

$ (vm) pip3 install ray[default]
$ (vm) ray start --head --port=6379 --num-cpus=0
...
# (Ray returns the IP address of the HEAD node, for example, RAY_HEAD_IP)
$ (vm) gcloud compute tpus tpu-vm create $TPU_NAME ... --metadata startup-script="pip3 install ray && ray start --address=$RAY_HEAD_IP --resources='{\"tpu_host\": 1}'"

Per praticità, forniamo anche script di base per creare una VM di coordinamento e implementare i contenuti di questa cartella nella VM di coordinamento. Per il codice sorgente, vedi create_cpu.sh e deploy.sh.

Questi script impostano alcuni valori predefiniti:

create_cpu.sh creerà una VM denominata $USER-admin e utilizzerà il progetto e la zona impostati per i valori predefiniti di gcloud config. Esegui gcloud config list per visualizzare i valori predefiniti.
Per impostazione predefinita, create_cpu.sh alloca un disco di avvio di 200 GB.
deploy.sh presuppone che il nome della VM sia $USER-admin. Se lo cambi in create_cpu.sh, assicurati di modificarlo anche in deploy.sh.

Per utilizzare gli script di utilità:

Clona il repository GitHub sulla tua macchina locale e vai alla ray_tpu:

$ git clone https://github.com/tensorflow/tpu.git
$ cd tpu/tools/ray_tpu/

Se non hai un account di servizio dedicato per l'amministrazione di TPU (vivamente consigliato), configurane uno:
```
$ ./create_tpu_service_account.sh
```
Nota: questo script deve essere eseguito una sola volta.
Crea una VM di coordinamento:
```
$ ./create_cpu.sh
```
Questo script installa le dipendenze sulla VM utilizzando un'istanza script di avvio e si blocca automaticamente fino al completamento dello script di avvio.
Esegui il deployment del codice locale nella VM di coordinamento:
```
$ ./deploy.sh
```
Accedi tramite SSH alla VM:
```
$ gcloud compute ssh $USER-admin -- -L8265:localhost:8265
```
Qui è abilitato il port forwarding, perché Ray avvierà automaticamente una dashboard la porta 8265. Dalla macchina da cui esegui SSH alla VM del coordinatore, potrai accedere a questa dashboard all'indirizzo http://127.0.0.1:8265/.
Se hai saltato il passaggio 0, configura le credenziali gcloud all'interno della VM CPU:
```
$ (vm) gcloud auth login --update-adc
```
Questo passaggio imposta le informazioni sull'ID progetto e consente l'esecuzione dell'API Cloud TPU con la VM coordinatore.

Attenzione: questo comando autorizza l'istanza VM a utilizzare il tuo Account Google personale, il che potrebbe rappresentare un rischio per la sicurezza in un ambiente di produzione.

Requisiti di installazione:

$ (vm) pip3 install -r src/requirements.txt

Avvia Ray sulla VM di coordinamento, che diventa il nodo principale del cluster Ray:
```
$ (vm) ray start --head --port=6379 --num-cpus=0
```
Nota: --num-cpus=0 eviterà che vengano pianificati job della CPU come la profilazione sul VM coordinatore.

Esempi di utilizzo

Esempio di JAX di base

run_basic_jax.py è un esempio minimo che dimostra come è possibile utilizzare i campi Ray Jobs e Ray un ambiente di runtime su un cluster Ray con VM TPU per eseguire un carico di lavoro JAX.

Per i framework ML compatibili con le Cloud TPU che utilizzano un modello di programmazione multi-controller, come JAX e PyTorch/XLA PJRT, devi eseguire almeno un processo per host. Per ulteriori informazioni, consulta Modello di programmazione multiprocesso. In pratica, il risultato potrebbe avere il seguente aspetto:

$ gcloud compute tpus tpu-vm scp my_bug_free_python_code my_tpu:~/ --worker=all
$ gcloud compute tpus tpu-vm ssh my_tpu --worker=all --command="python3 ~/my_bug_free_python_code/main.py"

Se hai più di circa 16 host, ad esempio una v4-128, incontrerai problemi di scalabilità di SSH e il comando potrebbe dover essere modificato in:

$ gcloud compute tpus tpu-vm scp my_bug_free_python_code my_tpu:~/ --worker=all --batch-size=8
$ gcloud compute tpus tpu-vm ssh my_tpu --worker=all --command="python3 ~/my_bug_free_python_code/main.py &" --batch-size=8

Questo può diventare un ostacolo sulla velocità degli sviluppatori se my_bug_free_python_code contiene insetti. Uno dei modi per risolvere questo problema è utilizzare un orchestratore come Kubernetes o Ray. Ray include il concetto di un Ambiente di runtime che, una volta applicato, esegue il deployment di codice e dipendenze quando l'applicazione Ray vengono eseguiti tutti i test delle unità.

La combinazione dell'ambiente di runtime di Ray con Ray Cluster e Ray Jobs ti consente di bypassare il ciclo SCP/SSH. Supponendo di aver seguito gli esempi precedenti, puoi eseguirlo con:

$ python3 legacy/run_basic_jax.py

L'output è simile al seguente:

2023-03-01 22:12:10,065   INFO worker.py:1364 -- Connecting to existing Ray cluster at address: 10.130.0.19:6379...
2023-03-01 22:12:10,072   INFO worker.py:1544 -- Connected to Ray cluster. View the dashboard at http://127.0.0.1:8265
W0301 22:12:11.148555 140341931026240 ray_tpu_controller.py:143] TPU is not found, create tpu...
Creating TPU:  $USER-ray-test
Request:  {'accelerator_config': {'topology': '2x2x2', 'type': 'V4'}, 'runtimeVersion': 'tpu-ubuntu2204-base', 'networkConfig': {'enableExternalIps': True}, 'metadata': {'startup-script': '#! /bin/bash\necho "hello world"\nmkdir -p /dev/shm\nsudo mount -t tmpfs -o size=100g tmpfs /dev/shm\n pip3 install ray[default]\nray start --resources=\'{"tpu_host": 1}\' --address=10.130.0.19:6379'}}
Create TPU operation still running...
...
Create TPU operation complete.
I0301 22:13:17.795493 140341931026240 ray_tpu_controller.py:121] Detected 0 TPU hosts in cluster, expecting 2 hosts in total
I0301 22:13:17.795823 140341931026240 ray_tpu_controller.py:160] Waiting for 30s for TPU hosts to join cluster...
…
I0301 22:15:17.986352 140341931026240 ray_tpu_controller.py:121] Detected 2 TPU hosts in cluster, expecting 2 hosts in total
I0301 22:15:17.986503 140341931026240 ray_tpu_controller.py:90] Ray already started on each host.
2023-03-01 22:15:18,010   INFO dashboard_sdk.py:315 -- Uploading package gcs://_ray_pkg_3599972ae38ce933.zip.
2023-03-01 22:15:18,010   INFO packaging.py:503 -- Creating a file package for local directory '/home/$USER/src'.
2023-03-01 22:15:18,080   INFO dashboard_sdk.py:362 -- Package gcs://_ray_pkg_3599972ae38ce933.zip already exists, skipping upload.
I0301 22:15:18.455581 140341931026240 ray_tpu_controller.py:169] Queued 2 jobs.
...
I0301 22:15:48.523541 140341931026240 ray_tpu_controller.py:254] [ADMIN]: raysubmit_WRUtVB7nMaRTgK39: Status is SUCCEEDED
I0301 22:15:48.561111 140341931026240 ray_tpu_controller.py:256] [raysubmit_WRUtVB7nMaRTgK39]: E0301 22:15:36.294834089   21286 credentials_generic.cc:35]            Could not get HOME environment variable.
8

I0301 22:15:58.575289 140341931026240 ray_tpu_controller.py:254] [ADMIN]: raysubmit_yPCPXHiFgaCK2rBY: Status is SUCCEEDED
I0301 22:15:58.584667 140341931026240 ray_tpu_controller.py:256] [raysubmit_yPCPXHiFgaCK2rBY]: E0301 22:15:35.720800499    8561 credentials_generic.cc:35]            Could not get HOME environment variable.
8

Addestramento con tolleranza di errore

Questo esempio mostra come utilizzare RayTpuController per implementare l'addestramento con tolleranza ai guasti. Per questo esempio, preaddestrando un semplice LLM PAX su una versione v4-16, ma tieni presente che puoi sostituire questo carico di lavoro PAX con qualsiasi altro carico di lavoro in esecuzione. Per il codice sorgente, vedi run_pax_autoresume.py

Per eseguire questo esempio:

Clona paxml nella VM del coordinatore:
```
$ git clone https://github.com/google/paxml.git
```
Per dimostrare la facilità d'uso offerta dall'ambiente di runtime Ray per apportare e implementare modifiche a JAX, questo esempio richiede di modificare PAX.

Aggiungi una nuova configurazione dell'esperimento:

$ cat <<EOT >> paxml/paxml/tasks/lm/params/lm_cloud.py

@experiment_registry.register
class TestModel(LmCloudSpmd2BLimitSteps):
ICI_MESH_SHAPE = [1, 4, 2]
CHECKPOINT_POLICY = layers.AutodiffCheckpointType.SAVE_CONTEXT_AND_OUT_PROJ

def task(self) -> tasks_lib.SingleTask.HParams:
  task_p = super().task()
  task_p.train.num_train_steps = 1000
  task_p.train.save_interval_steps = 100
  return task_p
EOT

Corsa di run_pax_autoresume.py:

$ python3 legacy/run_pax_autoresume.py --model_dir=gs://your/gcs/bucket

Durante l'esecuzione del carico di lavoro, esegui esperimenti su cosa succede quando elimini la TPU, per impostazione predefinita denominata $USER-tpu-ray:

$ gcloud compute tpus tpu-vm delete -q $USER-tpu-ray --zone=us-central2-b

Ray rileverà che la TPU non è disponibile con il seguente messaggio:

I0303 05:12:47.384248 140280737294144 checkpointer.py:64] Saving item to gs://$USER-us-central2/pax/v4-16-autoresume-test/checkpoints/checkpoint_00000200/metadata.
W0303 05:15:17.707648 140051311609600 ray_tpu_controller.py:127] TPU is not found, create tpu...
2023-03-03 05:15:30,774 WARNING worker.py:1866 -- The node with node id: 9426f44574cce4866be798cfed308f2d3e21ba69487d422872cdd6e3 and address: 10.130.0.113 and node name: 10.130.0.113 has been marked dead because the detector has missed too many heartbeats from it. This can happen when a       (1) raylet crashes unexpectedly (OOM, preempted node, etc.)
      (2) raylet has lagging heartbeats due to slow network or busy workload.
2023-03-03 05:15:33,243 WARNING worker.py:1866 -- The node with node id: 214f5e4656d1ef48f99148ddde46448253fe18672534467ee94b02ba and address: 10.130.0.114 and node name: 10.130.0.114 has been marked dead because the detector has missed too many heartbeats from it. This can happen when a       (1) raylet crashes unexpectedly (OOM, preempted node, etc.)
      (2) raylet has lagging heartbeats due to slow network or busy workload.

Il job ricrea automaticamente la VM TPU e riavvierà il job di addestramento in modo che possa riprendere l'addestramento dal checkpoint più recente (200 passaggi esempio):

I0303 05:22:43.141277 140226398705472 train.py:1149] Training loop starting...
I0303 05:22:43.141381 140226398705472 summary_utils.py:267] Opening SummaryWriter `gs://$USER-us-central2/pax/v4-16-autoresume-test/summaries/train`...
I0303 05:22:43.353654 140226398705472 summary_utils.py:267] Opening SummaryWriter `gs://$USER-us-central2/pax/v4-16-autoresume-test/summaries/eval_train`...
I0303 05:22:44.008952 140226398705472 py_utils.py:350] Starting sync_global_devices Start training loop from step: 200 across 8 devices globally

Ricerca degli iperparametri

Questo esempio mostra come utilizzare Ray Tune da Ray AIR per ottimizzare gli iperparametri MNIST di JAX/FLAX. Per il codice sorgente, vedi run_hp_search.py

Per eseguire questo esempio:

Installa i requisiti:

$ pip3 install -r src/tune/requirements.txt

Corsa di run_hp_search.py:

$ python3 src/tune/run_hp_search.py

L'output è simile al seguente:

Number of trials: 3/3 (3 TERMINATED)
+-----------------------------+------------+-------------------+-----------------+------------+--------+--------+------------------+
| Trial name                  | status     | loc               |   learning_rate |   momentum |    acc |   iter |   total time (s) |
|-----------------------------+------------+-------------------+-----------------+------------+--------+--------+------------------|
| hp_search_mnist_8cbbb_00000 | TERMINATED | 10.130.0.84:21340 |     1.15258e-09 |   0.897988 | 0.0982 |      3 |          82.4525 |
| hp_search_mnist_8cbbb_00001 | TERMINATED | 10.130.0.84:21340 |     0.000219523 |   0.825463 | 0.1009 |      3 |          73.1168 |
| hp_search_mnist_8cbbb_00002 | TERMINATED | 10.130.0.84:21340 |     1.08035e-08 |   0.660416 | 0.098  |      3 |          71.6813 |
+-----------------------------+------------+-------------------+-----------------+------------+--------+--------+------------------+

2023-03-02 21:50:47,378   INFO tune.py:798 -- Total run time: 318.07 seconds (318.01 seconds for the tuning loop).
...

Risoluzione dei problemi

Impossibile connettersi al nodo head Ray

Se esegui un carico di lavoro che crea/elimina il ciclo di vita delle TPU, a volte ciò non disconnette gli host TPU dal cluster Ray. Questo può vengono visualizzati come errori gRPC, che indicano che il nodo head Ray non è in grado di connettersi a un insieme di indirizzi IP.

Di conseguenza, potresti dover terminare la sessione di Ray (ray stop) e riavviarla (ray start --head --port=6379 --num-cpus=0).

Il job Ray non va a buon fine direttamente senza alcun output del log

PAX è sperimentale e questo esempio potrebbe non funzionare a causa delle dipendenze di pip. In questo caso, le informazioni potrebbero essere simili alle seguenti:

I0303 20:50:36.084963 140306486654720 ray_tpu_controller.py:174] Queued 2 jobs.
I0303 20:50:36.136786 140306486654720 ray_tpu_controller.py:238] Requested to clean up 1 stale jobs from previous failures.
I0303 20:50:36.148653 140306486654720 ray_tpu_controller.py:253] Job status: Counter({<JobStatus.FAILED: 'FAILED'>: 2})
I0303 20:51:38.582798 140306486654720 ray_tpu_controller.py:126] Detected 2 TPU hosts in cluster, expecting 2 hosts in total
W0303 20:51:38.589029 140306486654720 ray_tpu_controller.py:196] Detected job raysubmit_8j85YLdHH9pPrmuz FAILED.
2023-03-03 20:51:38,641   INFO dashboard_sdk.py:362 -- Package gcs://_ray_pkg_ae3cacd575e24531.zip already exists, skipping upload.
2023-03-03 20:51:38,706   INFO dashboard_sdk.py:362 -- Package gcs://_ray_pkg_ae3cacd575e24531.zip already exists, skipping upload.

Per visualizzare la causa principale dell'errore, puoi andare all'indirizzo http://127.0.0.1:8265/ e visualizzare la dashboard dei job in esecuzione/non riusciti, che fornirà maggiori informazioni. runtime_env_agent.log mostra tutte le informazioni sugli errori relativi alla configurazione di runtime_env, ad esempio:

60    INFO: pip is looking at multiple versions of  to determine which version is compatible with other requirements. This could take a while.
61    INFO: pip is looking at multiple versions of orbax to determine which version is compatible with other requirements. This could take a while.
62    ERROR: Cannot install paxml because these package versions have conflicting dependencies.
63
64    The conflict is caused by:
65        praxis 0.3.0 depends on t5x
66        praxis 0.2.1 depends on t5x
67        praxis 0.2.0 depends on t5x
68        praxis 0.1 depends on t5x
69
70    To fix this you could try to:
71    1. loosen the range of package versions you've specified
72    2. remove package versions to allow pip attempt to solve the dependency conflict
73
74    ERROR: ResolutionImpossible: for help visit https://pip.pypa.io/en/latest/topics/dependency-resolution/#dealing-with-dependency-conflicts