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.

• Proses debug performa PyTorch XLA pada VM TPU - bagian 1
• Proses debug performa PyTorch XLA pada VM TPU - bagian 2
• Proses debug performa PyTorch XLA pada 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
   

   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

Catatan: v4-8 menyertakan 4 perangkat TPU. Setiap pekerja hanya membuat profil satu perangkat TPU dari 4 perangkat yang tersedia. Pembuatan profil keempat perangkat TPU untuk setiap worker 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 di 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 mencakup langkah-langkah berikut.

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

Server dapat dimulai tepat di awal fungsi main.

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

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

Penting: Hindari menggabungkan xm.mark_step() dengan xp.Trace. Dalam hal ini, kode akan error dan menampilkan pesan error yang mirip dengan: "RuntimeError: Expecting scope to be empty but it is train_loop." xp.StepTrace tidak akan mengalami masalah ini karena 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 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.

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

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

Catatan: Penting untuk memiliki file .xplane dalam struktur bertingkat yang sama yang dihasilkan 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 terjadi 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.