torch.compile untuk inferensi

torch.compile mengompilasi kode PyTorch menjadi kernel yang dioptimalkan untuk mempercepat inferensi secara signifikan. Fitur ini bergantung pada TorchDynamo untuk mengompilasi kode menjadi graf dan TorchInductor untuk selanjutnya mengompilasi graf tersebut menjadi kernel yang dioptimalkan. Ini adalah alat optimasi yang sangat ampuh, dan dalam banyak kasus, hanya membutuhkan penambahan satu baris kode.

Bungkus model dengan torch.compile untuk mengompilasi dan mengembalikan model yang dioptimalkan.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model)

Panggilan awal ke torch.compile lambat karena model perlu dikompilasi. Panggilan berikutnya ke model yang sudah dikompilasi jauh lebih cepat karena tidak perlu dikompilasi lagi.

Terdapat beberapa parameter untuk menyesuaikan proses kompilasi. Dua di antaranya yang lebih penting tercantum di bawah ini. Untuk daftar lengkap parameter, silakan merujuk ke dokumentasi torch.compile.

Mode

Parameter mode menyediakan beberapa opsi performa untuk kompilasi. Coba berbagai mode untuk melihat mana yang paling sesuai dengan kasus penggunaan Anda.

• default adalah opsi yang seimbang antara kecepatan dan memori.
• reduce-overhead mengurangi overhead Python dengan mengorbankan sedikit lebih banyak memori, tetapi bisa lebih cepat.
• max-autotune menawarkan kecepatan tercepat, tetapi kompilasi memakan waktu lebih lama.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model, mode="reduce-overhead")

Fullgraph

Fullgraph berusaha mengompilasi seluruh model menjadi satu graf untuk memaksimalkan kinerja. torch.compile akan memicu error jika menemukan graph break, yang berarti model tidak dapat dikompilasi menjadi satu graf.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model, mode="reduce-overhead", fullgraph=True)

Benchmarks

Lihat tabel di bawah ini untuk benchmark kinerja yang membandingkan waktu inferensi rata-rata dalam milidetik dengan torch.compile diaktifkan dan dinonaktifkan di berbagai GPU dan ukuran batch pada gambar yang sama untuk berbagai tugas visi.

Pilih Subset pada tabel di bawah untuk beralih antara GPU yang berbeda, serta benchmark pada PyTorch nightly 2.1.0dev dan torch.compile dengan mode reduce-overhead diaktifkan.

Langkah selanjutnya

• Lihat panduan torch.compile untuk pelatihan untuk mempelajari cara menggunakan torch.compile untuk pelatihan.