Membuat profil workload PyTorch XLA

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

• Proses debug performa PyTorch XLA di VM TPU - bagian 1
• Proses debug performa PyTorch XLA di VM TPU - bagian 2
• Proses debug performa PyTorch XLA di VM TPU - bagian 3

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

   Mengekspor 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. Default ditampilkan dalam variabel yang diekspor, tetapi Anda juga dapat menggunakan salah satu 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

Catatan: v4-8 mencakup 4 perangkat TPU. Setiap pekerja hanya membuat profil satu perangkat TPU dari 4 perangkat yang tersedia. Membuat profil untuk keempat perangkat TPU bagi setiap pekerja akan segera didukung.

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 dalam skrip pelatihan. Mulai server dengan nomor port pilihan Anda, misalnya 9012 seperti yang ditunjukkan dalam perintah berikut.

Catatan: Anda dapat melihat contoh skrip ini yang menyertakan langkah-langkah berikut.

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

Server dapat dimulai tepat di awal fungsi main Anda.

Catatan: Jika Anda menggunakan modul Lightning, bukan memulai server, teruskan profiler="xla" ke trainer, dan server akan otomatis dimulai di port 9012.

Sekarang Anda dapat mengambil profil seperti yang dijelaskan di bagian berikut. Skrip ini membuat profil semua yang terjadi di satu perangkat TPU.

Menambahkan Trace

Jika Anda juga ingin membuat profil operasi di mesin host, Anda dapat menambahkan xp.StepTrace atau xp.Trace dalam kode. Fungsi ini melacak kode Python di mesin host. (Anda dapat menganggapnya sebagai pengukuran waktu yang diperlukan untuk mengeksekusi kode Python di 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 setiap bagian kode dengan

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

Penting: Hindari menggabungkan xm.mark_step() dengan xp.Trace. Dalam hal ini, kode akan error dengan pesan error yang mirip dengan: "RuntimeError: Expecting scope to be empty but it is train_loop." xp.StepTrace tidak akan mengalami masalah ini karena kode tersebut otomatis menambahkan xm.mark_step() saat keluar dari cakupan. Selain itu, Anda mungkin mendapatkan error ini saat menggunakan torchdynamo karena torchdynamo dapat memanggil xm.mark_step secara internal. Solusi saat ini adalah tidak menggunakan xp.Trace.

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

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

Perekaman 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 menyalinnya ke direktori beranda.

$ 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 hal 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 pilihan Anda. Anda juga dapat menyimpan langsung di bucket Cloud Storage. Untuk melakukannya, tetapkan logdir menjadi gs://your_bucket_name/.

Perekaman Terprogram

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

Catatan: trace_detached API hanya tersedia dalam build harian atau setelah rilis 2.2. Untuk mengambil gambar secara terprogram di versi torch_xla sebelumnya, Anda dapat menyalin isi metode untuk membuat thread pengambilan latar belakang secara langsung.

Misalnya, untuk mengambil profil secara otomatis 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 mesin yang sama atau di mesin lain (direkomendasikan).

Untuk menjalankan TensorBoard di komputer 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 di komputer tersebut di bagian ~/profiles/). Jika Anda menyimpan profil di direktori lain atau bucket Cloud Storage, pastikan untuk menentukan jalur dengan benar, misalnya, gs://your_bucket_name/profiles.

(vm)$ pip install tensorflow-cpu tensorboard-plugin-profile 

(vm)$ tensorboard --logdir ~/profiles/ --port 6006

Catatan: Anda harus memiliki file .xplane dalam struktur bertingkat yang sama yang dibuat saat Anda mengambil profil. Misalnya: ~/profiles/plugins/profile/2023_04_10_21_40_22/localhost_9012.xplane.pb

Catatan: Jika Anda mendapatkan error "plugin duplikat", hal ini mungkin karena ada beberapa paket dengan TensorBoard. Uninstal semua paket tersebut dan jalankan perintah pip install lagi. Misalnya, Anda mungkin perlu menjalankan:

(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.