torch.compile

Topik ini membahas torch.compile, yaitu lapisan kompilasi PyTorch yang mengubah kode model menjadi versi yang berjalan lebih cepat. Pembahasan mencakup cara menggunakannya pada model kecil, biaya panggilan pertama, penyebab graph break, pilihan mode kompilasi, dan kondisi ketika kompilasi tidak memberi keuntungan. Setelah topik ini dibaca, model yang sama dapat dijalankan dalam dua bentuk, yaitu versi eager dan versi hasil kompilasi, disertai pemahaman tentang kapan masing-masing bentuk menguntungkan.

Apa yang dilakukan torch.compile#

Secara bawaan, PyTorch menjalankan model secara eager, yakni setiap operasi dieksekusi satu per satu segera ketika Python membaca barisnya. Cara ini fleksibel dan mudah di-debug, tetapi menyisakan overhead Python pada setiap panggilan operasi, yaitu waktu yang terpakai untuk menafsirkan kode dan menyiapkan kernel alih-alih menghitung. torch.compile memangkas overhead tersebut dengan menganalisis fungsi forward secara keseluruhan sebelum menjalankannya.

Prosesnya berlangsung dalam dua tahap. Tahap pertama dilakukan oleh dynamo, komponen yang membaca kode Python pada forward dan merekamnya menjadi satu graf komputasi yang utuh. Tahap kedua dilakukan oleh inductor, komponen yang menerjemahkan graf tersebut menjadi kode C++ yang efisien, mengompilasinya dengan gcc, lalu menjalankan hasil kompilasi itu. Salah satu penghematannya datang dari fusion, yaitu penggabungan beberapa operasi kecil menjadi satu kernel sehingga tensor perantara tidak perlu ditulis ke memori lalu dibaca ulang.

Penggunaannya berupa satu baris. Model dibungkus dengan torch.compile, dan hasil pembungkusannya digunakan seperti model biasa. Objek aslinya tidak diubah, sehingga bobot yang sama dapat diakses dari kedua sisi.

Penggunaan dasar#

Contoh berikut mendefinisikan jaringan kecil dengan dua layer Linear, lalu menjalankannya dalam dua bentuk pada masukan yang sama. Satu-satunya argumen torch.compile adalah objek model, dan hasilnya disimpan pada variabel baru, sehingga kode di sekitarnya tidak berubah. Nilai acak diberi seed terlebih dahulu agar masukan yang digunakan kedua bentuk identik dan perbandingannya adil. Perhatikan bahwa panggilan pertama terhadap model_kompilasi memicu proses kompilasi, sehingga baris tersebut membutuhkan waktu lebih lama dibandingkan dengan baris lainnya.

import torch
import torch.nn as nn

class JaringanKecil(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(32, 64)
        self.fc2 = nn.Linear(64, 10)

    def forward(self, x):
        return self.fc2(torch.relu(self.fc1(x)))

torch.manual_seed(42)
model = JaringanKecil()
model_kompilasi = torch.compile(model)
print("Tipe hasil torch.compile :", type(model_kompilasi).__name__)

x = torch.randn(8, 32)
y_eager = model(x)
y_kompilasi = model_kompilasi(x)  # panggilan pertama: kompilasi berlangsung di sini
print("Keluaran sama (atol 1e-5):", torch.allclose(y_eager, y_kompilasi, atol=1e-5))

Tipe hasil torch.compile : OptimizedModule
Keluaran sama (atol 1e-5): True

torch.compile mengembalikan objek bertipe OptimizedModule, yaitu pembungkus yang meneruskan setiap panggilan ke kode hasil kompilasi. Keluaran kedua menunjukkan bahwa kedua bentuk memberikan hasil yang sama pada toleransi 1e-5. Perbedaan numerik yang sangat kecil memang dapat muncul karena urutan operasi pada versi hasil kompilasi tidak selalu identik dengan versi eager, sehingga perbandingan hasil selalu menggunakan toleransi dan bukan kesamaan bit. Model aslinya pun tetap dapat digunakan seperti biasa, karena pembungkus tersebut tidak mengubah objek model yang lama.

Biaya panggilan pertama dan warm-up#

Kompilasi tidak berlangsung saat torch.compile dipanggil, melainkan saat pemanggilan model yang pertama. Pada saat itu, dynamo merekam graf, inductor menghasilkan kode, dan gcc mengompilasinya. Pada model kecil di atas, pengukuran yang dilakukan saat materi ini ditulis menunjukkan panggilan pertama memakan waktu sekitar delapan detik pada CPU, sedangkan panggilan berikutnya turun ke orde milidetik. Sebagian besar hasil kompilasi disimpan dalam cache di disk, sehingga program yang dijalankan ulang tidak selalu mengulang seluruh prosesnya.

Pola pengukuran yang benar mengikuti fakta tersebut. Berikut adalah kerangka pengukuran yang membandingkan kedua bentuk setelah warm-up, yaitu beberapa panggilan pembuka yang memastikan kompilasi dan kernel sudah siap:

# tidak-dijalankan
# angka waktu berubah antar mesin dan antar eksekusi, sehingga potongan ini
# tidak diverifikasi ulang secara ketat. Angka pada prosa berasal dari
# pengukuran saat materi ini ditulis
for _ in range(10):
    model(x)                      # warm-up
t0 = time.perf_counter()
for _ in range(200):
    model(x)
t_eager = time.perf_counter() - t0

model_kompilasi(x)                # panggilan pertama: kompilasi berlangsung di sini
for _ in range(10):
    model_kompilasi(x)            # warm-up
t0 = time.perf_counter()
for _ in range(200):
    model_kompilasi(x)
t_compiled = time.perf_counter() - t0

Pada pengukuran tersebut, versi hasil kompilasi berjalan sekitar dua setengah kali lebih cepat daripada versi eager untuk model di atas. Angka pasti bergantung pada arsitektur model, bentuk masukan, dan mesin yang digunakan, sehingga yang penting untuk diikuti adalah polanya: warm-up dahulu, ukur setelahnya.

Graph break#

Dynamo tidak selalu dapat merekam seluruh fungsi forward menjadi satu graf. Operasi yang nilainya baru diketahui saat berjalan, seperti percabangan yang bergantung pada isi tensor, efek samping seperti print, atau pemanggilan pustaka di luar PyTorch, membuat dynamo membagi graf menjadi beberapa bagian. Peristiwa ini disebut graph break, yaitu titik ketika perekaman terputus dan eksekusi kembali ke Python di antara dua graf.

Fungsi torch._dynamo.explain menganalisis model tanpa menjalankan kompilasi penuh, lalu melaporkan jumlah graf yang dihasilkan beserta banyaknya titik putus:

import torch
import torch.nn as nn

class JaringanPutus(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(32, 64)
        self.fc2 = nn.Linear(64, 10)

    def forward(self, x):
        y = torch.relu(self.fc1(x))
        if y.sum() > 0:  # cabang bergantung pada nilai tensor
            y = y * 1.0
        return self.fc2(y)

torch.manual_seed(0)
model = JaringanPutus()
x = torch.randn(4, 32)
laporan = torch._dynamo.explain(model)(x)
print("Jumlah graf   :", laporan.graph_count)
print("Graph break   :", laporan.graph_break_count)
print("Jumlah operasi:", laporan.op_count)

Jumlah graf   : 2
Graph break   : 1
Jumlah operasi: 5

Keluaran tersebut dapat dibaca langsung. Satu percabangan yang bergantung pada y.sum() memecah forward menjadi dua graf dengan satu graph break di antara keduanya, dan hanya lima operasi yang masuk ke dalam graf. Bagian forward yang berada di antara dua graf tetap dijalankan oleh Python, sehingga keuntungan kompilasinya mengecil. Model dengan banyak graph break masih berjalan benar, tetapi hasil kompilasinya mendekati versi eager.

Mode kompilasi#

Parameter mode mengatur seberapa jauh inductor bekerja untuk mencari versi tercepat. Mode default sudah memadai untuk sebagian besar kasus, sedangkan dua mode lainnya menawarkan pertukaran antara waktu kompilasi dan kecepatan hasil akhir.

Penggunaannya tetap satu baris, misalnya torch.compile(model, mode="reduce-overhead"). Mode max-autotune memakan waktu kompilasi paling lama karena banyak percobaan dilakukan saat kompilasi, dan mode reduce-overhead memberi manfaat terbesar pada GPU karena CUDA graphs tidak berlaku pada CPU. Pemilihan mode tidak mengubah hasil perhitungan, hanya kecepatan dan lama kompilasinya. Untuk model kecil pada CPU sebagaimana digunakan sepanjang jalur ini, mode default biasanya sudah cukup, dan pengukuran sebaiknya dilakukan sebelum memutuskan menggunakan mode yang lebih mahal.

Kapan kompilasi tidak membantu#

Kompilasi adalah investasi, sehingga ada kondisi ketika biayanya tidak kembali. Kondisi pertama adalah model yang sangat kecil atau sangat jarang dipanggil. Waktu kompilasi tetap harus dibayar, sementara penghematan per panggilan kecil, sehingga total waktu justru bertambah. Kondisi kedua adalah model dengan banyak graph break, sebagaimana ditunjukkan pada bagian sebelumnya, karena bagian yang ditinggalkan kepada Python justru bagian yang semula menjadi sasaran pengoptimalan.

Kondisi ketiga adalah bentuk masukan yang terus berubah. Kompilasi mengikat graf pada bentuk masukan yang direkam, sehingga batch berukuran 8 dan batch berukuran 9 memicu dua kompilasi berbeda. Dynamo menyimpan sejumlah versi dan melakukan kompilasi ulang ketika bentuk baru muncul, dengan batas yang ditetapkan konfigurasi. Layanan yang menerima batch dengan ukuran bebas perlu mempertimbangkan hal ini sebelum mengaktifkan kompilasi.

Ringkasan#

• torch.compile membungkus model menjadi OptimizedModule: dynamo merekam forward menjadi graf, dan inductor menghasilkan kode C++ hasil kompilasi yang memangkas overhead Python melalui fusion.
• Biaya kompilasi muncul pada panggilan pertama dan sebagian disimpan dalam cache, sehingga pengukuran performa selalu dilakukan setelah warm-up.
• Operasi yang bergantung pada nilai tensor, efek samping, atau pustaka luar menyebabkan graph break, yaitu pemisahan graf yang mengurangi keuntungan kompilasi.
• Mode default, reduce-overhead, dan max-autotune mengatur pertukaran antara lama kompilasi dan kecepatan, sedangkan model kecil, banyak graph break, dan bentuk masukan yang terus berubah membuat kompilasi tidak menguntungkan.

Referensi#

• torch.compiler — PyTorch documentation — Dokumentasi lapisan compiler, termasuk penjelasan dynamo, inductor, dan graph break.
• Introduction to torch.compile — PyTorch Tutorials — Tutorial resmi yang mendemonstrasikan penggunaan torch.compile dan graph break.
• torch.compile — PyTorch documentation — Rujukan parameter fungsi torch.compile beserta mode-mode bawaannya.