Panduan performa Cloud TPU

Langkah pertama saat memecahkan masalah performa TPU adalah membuat profil model Anda. Untuk mengetahui informasi selengkapnya tentang merekam profil performa, lihat Membuat profil model Anda di Cloud TPU.

Performa model TPU

Bagian ini menjelaskan masalah umum yang dapat mengurangi performa model dan cara mengatasinya.

1. Model terikat input

   TPU melakukan perhitungan dengan sangat cepat. Untuk memastikan TPU tidak dalam kondisi idle, Anda harus memastikan ada aliran data yang stabil yang dimuat ke TPU. Cara melakukannya bergantung pada cara Anda memuat dan memproses awal set data. Misalnya, Anda dapat membaca file data secara paralel menggunakan tf.data.TFRecordset() dan parameter num_parallel_reads.

2. Ukuran batch terlalu kecil karena sharding (membagi batch di seluruh core)

   Runtime TPU membagi batch di semua 8 core perangkat TPU (misalnya v2-8 atau v3-8). Jika Anda menentukan ukuran batch global sebesar 128, setiap core akan menerima ukuran batch sebesar 16 (128 / 8).

   Untuk penggunaan memori yang optimal, gunakan ukuran batch terbesar yang sesuai dengan memori TPU. Setiap core TPU menggunakan register vektor 8 X 128 dua dimensi untuk memproses perkalian matriks. Secara umum, ukuran batch Anda harus dapat dibagi rata dengan 8 atau 128.

3. Penyesuaian Pengelolaan Memori

   Anda dapat menggunakan variabel lingkungan TPU_PREMAPPED_BUFFER_SIZE untuk menyetel perilaku runtime tingkat rendah.

• Deskripsi: TPU_PREMAPPED_BUFFER_SIZE menetapkan ukuran buffer memori host (dalam byte) yang dipetakan sebelumnya dan disematkan untuk digunakan oleh runtime TPU untuk transfer data (misalnya, DMA). Nilai defaultnya adalah 4294967296 byte. Nilai harus kelipatan 2^12 (4 KB = 4 * 1024 Byte = 4096 = 2^12).

  Contoh berikut adalah nilai TPU_PRE_MAPPED_BUFFER_SIZE yang valid.

      17179869184 = 2^34 = 2^22 * 2^12 (2^22 4KB pages will be premapped).
      40000000000 = 5^10 * 2^12 = (5^10 4KB pages will be premapped).
  

• Dampak: Meningkatkan ukuran ini berpotensi meningkatkan performa transfer data antara host dan perangkat TPU, terutama untuk beban kerja dengan tensor besar atau komunikasi host-perangkat yang sering. Namun, tindakan ini juga meningkatkan jumlah memori host yang disematkan, sehingga mengurangi memori yang tersedia untuk proses lain.

  Ukuran buffer

  Jika region buffer yang dipetakan sebelumnya tidak cukup besar untuk mengalokasikan memori selama runtime program, workload akan gagal dan menampilkan error RESOURCE_EXHAUSTED yang mirip dengan:

  "Mengalokasikan buffer dari region yang telah dipetakan sebelumnya gagal dengan: RESOURCE_EXHAUSTED: Mencoba mengalokasikan allocation_size. Hal itu tidak mungkin dilakukan. Ada available_size gratis."

  Jika buffer terlalu besar, inisialisasi TPU dapat memerlukan waktu lebih lama (mungkin lebih dari 15 detik), sehingga TPU tampak macet.

  Untuk mendiagnosisnya, periksa log runtime TPU. Log ini menjelaskan operasi yang sedang dilakukan, termasuk pra-pemetaan buffer. Anda dapat menemukan log di /tmp/tpu_logs/tpu_driver.INFO atau mencetaknya langsung ke konsol dengan menetapkan variabel lingkungan TPU_STDERR_LOG_LEVEL=0. Setelan ini akan menghasilkan output yang mirip dengan:

   I0604 12:45:24.926233   62136 tpu_hal.cc:214] Starting premapped memory manager initialization...
   I0604 12:45:29.411218   62136 system.cc:1059] tpu::System initialized, current host id: 0, logical device ids: 0
   I0604 12:45:29.411244   61600 tfrt_tpu_system_state.cc:216] CreateTpuSystemState: TPU initialization is successful and it took 5.583190661s
   I0604 12:45:29.411267   61600 tfrt_tpu_system_state.cc:220] CreateTpuSystemState: using TPU host premapped buffer of size: 4294967296
   ```
  
  This output will tell you how long it took to initialize the TPU and
  the size of the premapped buffer.
  

• Penggunaan: Jika buffer yang dipetakan sebelumnya terlalu kecil atau terlalu besar, Anda dapat menetapkan ukuran buffer secara manual menggunakan variabel lingkungan berikut.

  TPU_PREMAPPED_BUFFER_SIZE: Sets the total size (in bytes) of the
  pre-mapped buffer region.
  TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES: Sets the maximum size of
  a single buffer that can be allocated from the pre-mapped region.
  

  Misalnya, Anda dapat:

   export TPU_PREMAPPED_BUFFER_SIZE=4294967296
  

  untuk menetapkan ukuran buffer dan:

   export TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES
   ```
   to enable it.
  
   This export sets the size to the default.
  

• Panduan: Sesuaikan nilai TPU_PREMAPPED_BUFFER_SIZE jika Anda menduga transfer data host-perangkat menjadi hambatan. Pantau penggunaan memori host dan performa model untuk menemukan keseimbangan yang optimal. Nilai default biasanya sudah cukup untuk sebagian besar kasus penggunaan.

Pengoptimalan compiler XLA

XLA adalah compiler untuk machine learning yang dapat menghasilkan biner untuk TPU, CPU, GPU, dan platform lainnya. Meskipun XLA adalah bagian dari basis kode TensorFlow standar, XLA juga dapat digunakan pada model PyTorch dan JAX. Model untuk Cloud TPU diterjemahkan ke grafik XLA, yang kemudian dikompilasi XLA menjadi file yang dapat dieksekusi TPU. Untuk mengetahui informasi selengkapnya tentang XLA, lihat XLA: Pengoptimalan Compiler untuk Machine Learning.

Padding

Untuk menggunakan memori TPU secara efisien, susun data Anda sehingga dapat dibagi menjadi potongan 128 x 8. Jika data untuk komputasi matriks tidak mengisi seluruh chunk 128 x 8, compiler XLA akan mengisi tensor. Ada dua kelemahan padding:

1. Tensor yang di-padding kurang memanfaatkan core TPU.
2. Padding meningkatkan jumlah penyimpanan memori on-chip yang diperlukan untuk tensor dan dapat menyebabkan error kehabisan memori.

Meskipun pengoptimalan dilakukan secara otomatis oleh compiler XLA jika diperlukan, Anda dapat menentukan jumlah pengoptimalan yang dilakukan menggunakan alat penampil memori. Anda dapat menghindari padding dengan memilih dimensi tensor yang cocok untuk TPU.

Dimensi tensor

Untuk mencapai FLOP puncak, dimensi perkalian matriks harus lebih besar daripada ukuran MXU untuk versi TPU yang Anda gunakan. Ukuran MXU adalah 256 x 256 untuk v6e dan 128 x 128 untuk versi sebelum v6e. Untuk mengetahui informasi selengkapnya, lihat Arsitektur sistem Cloud TPU.

Ukuran batch

Compiler XLA membulatkan ukuran tensor yang disimpan dalam memori HBM TPU untuk melakukan komputasi secara lebih efisien. Pengisihan ini terjadi secara transparan di tingkat hardware dan tidak memengaruhi hasil. Namun, dalam kasus tertentu, padding dapat menyebabkan peningkatan penggunaan memori dan waktu eksekusi secara signifikan.

Runtime TPU menata tensor dalam memori untuk memaksimalkan efisiensi komputasi dan meminimalkan padding. Untuk meminimalkan overhead memori dan memaksimalkan efisiensi komputasi, salah satu kondisi berikut harus terpenuhi:

1. Ukuran batch total harus kelipatan 64 (8 per core TPU), dan ukuran dimensi fitur harus kelipatan 128.

2. Ukuran batch total harus kelipatan 1024 (128 per core TPU), dan ukuran dimensi fitur harus kelipatan 8.

Menggunakan ukuran batch 1024 dan dimensi fitur yang merupakan kelipatan 128 akan menghasilkan efisiensi terbaik, meskipun hal ini mungkin tidak dapat dilakukan untuk semua model.

Catatan: Dimensi fitur mengacu pada ukuran tersembunyi dari lapisan yang terhubung sepenuhnya atau jumlah saluran output dalam konvolusi. Tidak semua lapisan dapat mematuhi aturan ini, terutama lapisan pertama dan terakhir jaringan. Tidak masalah, sebagian besar model memerlukan sejumlah padding.

Fusion

Fusion adalah teknik umum yang digunakan compiler XLA untuk mengoptimalkan program. Operasi gabungan adalah kombinasi dari beberapa operasi konstituen yang akan dieksekusi bersama-sama.

Misalnya, pertimbangkan serangkaian operasi berikut:

    tmp = tf.add(x, y)
    result = tf.multiply(tmp, z)

Kode ini kira-kira setara dengan kode semu berikut:

    for (i = 0; i < element_count; i++) {
      tmp[i] = x[i] + y[i];
    }

    for (i = 0; i < element_count; i++) {
      result[i] = tmp[i] * z[i];
    }

Dengan penggabungan, akses array terjadi secara bersamaan:

    for (i = 0; i < element_count; i++) {
      result[i] = (x[i] + y[i]) * z[i];
    }

Dalam contoh ini, jumlah perjalanan pulang pergi memori berkurang dan XLA tidak perlu mengalokasikan ruang untuk 'tmp'.

Penggabungan adalah pengoptimalan penting dan memberikan manfaat bagi Cloud TPU dalam beberapa cara:

• Hal ini mengurangi transfer memori dengan menghilangkan kebutuhan untuk menyimpan hasil perantara dalam memori utama, yang lambat.
• Hal ini memungkinkan pemanfaatan unit hardware yang lebih besar yang jika tidak, akan tidak digunakan.
• Hal ini dapat mengurangi pemanfaatan memori model karena lebih sedikit buffer yang perlu aktif secara bersamaan.

Menyiarkan

Penyiaran terjadi secara implisit saat dua tensor dengan bentuk yang berbeda, tetapi kompatibel, digabungkan.

Misalnya, tf.add(vector, matrix) mengharuskan vektor disiarkan ke bentuk matriks. Hasil operasi memiliki bentuk yang sama dengan matriks. Untuk mengetahui detail selengkapnya, lihat panduan untuk penyiaran array.

Meskipun siaran sering kali dapat digabungkan dengan konsumennya, memaksa siaran dapat mengakibatkan performa yang buruk dan peningkatan penggunaan memori.

Dalam contoh berikut, siaran implisit dalam penambahan vektor dan matriks tidak dapat digabungkan dengan argmax yang menghasilkan siaran yang diwujudkan:

`tf.argmax(tf.add(vector, zero_matrix), axis=0)`