Membuat profil workload PyTorch XLA

Pembuatan profil adalah cara untuk menganalisis dan meningkatkan performa model. Meskipun ada lebih banyak hal, terkadang ada baiknya menganggap pembuatan profil sebagai operasi pengaturan waktu dan bagian dari kode yang berjalan di kedua perangkat (TPU) dan host (CPU). Panduan ini memberikan ringkasan singkat tentang cara membuat profil kode untuk pelatihan atau inferensi. Untuk mengetahui informasi selengkapnya tentang cara menganalisis profil yang dihasilkan, lihat panduan berikut.

Mulai

Membuat TPU

  1. Mengekspor variabel lingkungan:

    $ export TPU_NAME=your_tpu_name
$ export ZONE=us-central2-b
$ export PROJECT_ID=project-id
$ export ACCELERATOR_TYPE=v4-8
$ export RUNTIME_VERSION=tpu-vm-v4-pt-2.0

    Ekspor deskripsi variabel

    TPU name
    Nama yang ingin Anda gunakan untuk Cloud TPU.
    zone
    Zona tempat Anda berencana membuat Cloud TPU.
    project ID
    Project ID yang Anda gunakan untuk melatih dan membuat profil model.
    accelerator-type
    Jenis akselerator menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Untuk mengetahui informasi selengkapnya tentang jenis akselerator yang didukung untuk setiap versi TPU, lihat versi TPU.
    version
    Versi runtime Cloud TPU. Nilai default ditampilkan dalam variabel yang diekspor, tetapi Anda juga dapat menggunakan salah satunya dari daftar konfigurasi yang didukung.

  2. Meluncurkan resource TPU

    $ gcloud compute tpus tpu-vm create ${TPU_NAME} \
--zone us-central2-b \
--accelerator-type ${ACCELERATOR_TYPE} \
--version ${RUNTIME_VERSION} \
--project $PROJECT_ID \
--subnetwork=tpusubnet

  3. Pindahkan kode Anda ke direktori utama di VM TPU menggunakan perintah gcloud scp. Contoh:

    $ gcloud compute tpus tpu-vm scp my-code-file ${TPU_NAME}: --zone ${ZONE}

Pembuatan profil

Profil dapat diambil secara manual melalui capture_profile.py atau secara terprogram dari dalam skrip pelatihan menggunakan torch_xla.debug.profiler API.

Memulai Server Profil

Untuk mengambil profil, server profil harus berjalan di dalam skrip pelatihan. Mulai server dengan nomor port pilihan Anda, misalnya 9012 seperti yang ditunjukkan dalam perintah berikut.

import torch_xla.debug.profiler as xp
server = xp.start_server(9012)

Server dapat dimulai tepat di awal fungsi main.

Kini Anda dapat mengambil profil seperti yang dijelaskan di bagian berikut. Skrip ini akan memprofilkan semua yang terjadi di satu perangkat TPU.

Menambahkan Trace

Jika juga ingin membuat profil operasi pada mesin host, Anda dapat menambahkan xp.StepTrace atau xp.Trace dalam kode Anda. Fungsi ini melacak kode Python di mesin host. (Anda dapat menganggapnya sebagai mengukur lama waktu yang diperlukan untuk mengeksekusi kode Python pada host (CPU) sebelum meneruskan "grafik" ke perangkat TPU. Jadi, ini sangat berguna untuk menganalisis overhead pelacakan). Anda dapat menambahkannya di dalam loop pelatihan tempat kode memproses batch data, misalnya,

for step, batch in enumerate(train_dataloader):
    with xp.StepTrace('Training_step', step_num=step):
       ...

atau menggabungkan bagian individual dari kode dengan

 with xp.Trace('loss'):
    loss = ...

Jika menggunakan Pencahayaan, Anda dapat melewati penambahan rekaman aktivitas karena dilakukan secara otomatis di beberapa bagian kode. Namun, jika ingin menambahkan trace tambahan, Anda dapat memasukkannya di dalam loop pelatihan.

Anda akan dapat merekam aktivitas perangkat setelah kompilasi awal; tunggu hingga model memulai langkah pelatihan atau inferensinya.

Ambil Manual

Skrip capture_profile.py dari repositori Pytorch XLA memungkinkan pengambilan profil dengan cepat. Anda dapat melakukannya dengan menyalin file profil pengambilan langsung ke VM TPU. Perintah berikut akan menyalinnya ke direktori {i>home<i}.

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--zone us-central2-b \
--worker=all \
--command="wget https://raw.githubusercontent.com/pytorch/xla/master/scripts/capture_profile.py"

Saat pelatihan berjalan, jalankan perintah berikut untuk mengambil profil:

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--zone us-central2-b \
--worker=all \
--command="python3 capture_profile.py --service_addr "localhost:9012" --logdir ~/profiles/ --duration_ms 2000"

Perintah ini menyimpan file .xplane.pb di logdir. Anda dapat mengubah direktori logging ~/profiles/ ke lokasi dan nama yang diinginkan. Anda juga dapat menyimpannya secara langsung di bucket Cloud Storage. Untuk melakukannya, tetapkan logdir ke gs://your_bucket_name/.

Perekaman Terprogram

Daripada mengambil profil secara manual dengan memicu skrip, Anda dapat mengonfigurasi skrip pelatihan untuk memicu profil secara otomatis menggunakan API torch_xla.debug.profiler.trace_detached dalam skrip pelatihan Anda.

Sebagai contoh, untuk secara otomatis mengambil profil pada epoch dan langkah tertentu, Anda dapat mengonfigurasi skrip pelatihan untuk menggunakan variabel lingkungan PROFILE_STEP, PROFILE_EPOCH, dan PROFILE_LOGDIR:

import os
import torch_xla.debug.profiler as xp

# Within the training script, read the step and epoch to profile from the
# environment.
profile_step = int(os.environ.get('PROFILE_STEP', -1))
profile_epoch = int(os.environ.get('PROFILE_EPOCH', -1))
...

for epoch in range(num_epoch):
   ...
   for step, data in enumerate(epoch_dataloader):
      if epoch == profile_epoch and step == profile_step:
         profile_logdir = os.environ['PROFILE_LOGDIR']
         # Use trace_detached to capture the profile from a background thread
         xp.trace_detached('localhost:9012', profile_logdir)
      ...

Tindakan ini akan menyimpan file .xplane.pb di direktori yang ditentukan oleh variabel lingkungan PROFILE_LOGDIR.

Analisis di TensorBoard

Untuk menganalisis profil lebih lanjut, Anda dapat menggunakan TensorBoard dengan plugin TPU TensorBoard
di komputer yang sama atau di komputer lain (direkomendasikan).

Untuk menjalankan TensorBoard pada mesin jarak jauh, hubungkan ke komputer tersebut menggunakan SSH dan aktifkan penerusan port. Misalnya,

$ ssh -L 6006:localhost:6006 remote server address

atau

$ gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --ssh-flag="-4 -L 6006:localhost:6006"

Di komputer jarak jauh, instal paket yang diperlukan dan luncurkan TensorBoard (dengan asumsi Anda memiliki profil pada mesin tersebut di ~/profiles/). Jika Anda menyimpan profil di direktori atau bucket Cloud Storage lain, pastikan untuk menentukan jalur dengan benar, misalnya, gs://your_bucket_name/profiles.

(vm)$ pip install tensorflow tensorboard-plugin-profile
 
(vm)$ tensorboard --logdir ~/profiles/ --port 6006
 
(vm)$ pip uninstall tensorflow tf-nightly tensorboard tb-nightly tbp-nightly

Menjalankan TensorBoard

Di browser lokal, buka: http://localhost:6006/ dan pilih PROFILE dari menu drop-down untuk memuat profil Anda.

Lihat alat TPU untuk mengetahui informasi tentang alat TensorBoard dan cara menafsirkan output.

gambar

gambar