Pelatihan TPU

TPU (Tensor Processing Unit) adalah jenis hardware yang dirancang khusus untuk melatih model secara efisien. Accelerate mendukung pelatihan TPU, tetapi ada beberapa hal yang perlu Anda perhatikan, yaitu kompilasi graf. Tutorial ini membahas kompilasi secara singkat, dan untuk detail lebih lanjut, lihat panduan Training on TPUs with Accelerate.

Kompilasi

TPU membuat graf dari semua operasi dalam langkah pelatihan, seperti forward pass, backward pass, dan langkah optimizer. Inilah mengapa langkah pelatihan pertama selalu membutuhkan waktu karena membangun dan mengompilasi graf ini memakan waktu. Namun, begitu kompilasi selesai, hasilnya di-cache dan semua langkah berikutnya jauh lebih cepat.

Kuncinya adalah menghindari kompilasi ulang kode Anda, karena jika tidak, proses pelatihan akan menjadi sangat lambat. Ini berarti semua operasi Anda harus persis sama:

• semua tensor dalam batch Anda harus memiliki panjang yang sama (misalnya, tidak ada dynamic padding untuk tugas NLP)
• kode Anda harus statis (misalnya, tidak ada lapisan dengan for loop yang memiliki panjang berbeda tergantung pada input seperti LSTM)

Pengikatan bobot

Desain model bahasa yang umum adalah mengaitkan bobot dari lapisan embedding dan softmax. Namun, memindahkan model ke TPU (baik secara mandiri maupun dengan melewatkan model ke metode prepare()) akan memutus pengaitan bobot tersebut, sehingga Anda perlu mengaitkan ulang bobotnya.

Untuk menambahkan perilaku khusus (seperti weight tying) di skrip Anda untuk TPU, atur distributed_type ke DistributedType.TPU terlebih dahulu. Kemudian Anda dapat menggunakan metode tie_weights untuk mengikat bobotnya.

if accelerator.distributed_type == DistributedType.TPU:
    model.tie_weights()