Kompilasi

Ikhtisar

Pytorch 2.0 memperkenalkan torch.compile, fitur yang ampuh untuk membuat kode PyTorch berjalan lebih cepat dengan mengompilasi kode PyTorch secara JIT menjadi kernel yang dioptimalkan. Fitur utama dari torch.compile meliputi:

• Perbaikan Kinerja: Mempercepat eksekusi model secara signifikan dengan mengoptimalkan graf komputasi.
• Kemudahan Penggunaan: Memerlukan perubahan kode minimal untuk diimplementasikan, sehingga sangat mudah diakses.
• Kompatibilitas: Bekerja mulus dengan kode dan model PyTorch yang ada.

Saat digunakan dengan Accelerate, torch.compile terintegrasi dengan mulus ke dalam alur kerja pelatihan terdistribusi, memungkinkan Anda untuk mendapatkan manfaat dari eksekusi terdistribusi dan optimasi kompilasi secara bersamaan.

Eksekusi pertama kode yang dikompilasi biasanya memakan waktu lebih lama karena mencakup waktu kompilasi, tetapi eksekusi berikutnya jauh lebih cepat. Untuk kinerja optimal dalam berbagai skenario, torch.compile menawarkan berbagai mode seperti "default", "reduce-overhead" (yang menggunakan CUDA graphs untuk mengurangi overhead lebih lanjut), dan "max-autotune" (yang melakukan autotuning ekstensif untuk menemukan kernel terbaik untuk model Anda).

Menggunakan torch.compile dengan Accelerate

Accelerate menyediakan TorchDynamoPlugin untuk integrasi torch.compile yang mudah dan mulus ke dalam skrip pelatihan Anda.

from accelerate import Accelerator
from accelerate.utils import TorchDynamoPlugin

# Configure the compilation backend
dynamo_plugin = TorchDynamoPlugin(
    backend="inductor",  # Options: "inductor", "aot_eager", "aot_nvfuser", etc.
    mode="default",      # Options: "default", "reduce-overhead", "max-autotune"
    fullgraph=True,
    dynamic=False
)

# Initialize accelerator with the plugin
accelerator = Accelerator(dynamo_plugin=dynamo_plugin)
# This will apply torch.compile to your model
model = accelerator.prepare(model)

Kompatibel dengan semua fitur dan plugin lain dari Accelerate, termasuk mixed precision, pelatihan terdistribusi (DDP, FSDP, Deepspeed), dan sebagainya.

Kompilasi Regional

Alih-alih mencoba mengompilasi seluruh model, yang biasanya memiliki ruang masalah yang besar untuk optimasi. Kompilasi regional menargetkan blok berulang dari kelas yang sama dan mengompilasinya secara berurutan untuk memanfaatkan cache compiler. Misalnya, dalam GPT2LMHeadModel, blok/kelas yang berulang adalah GPT2Block, dan dapat diakses sebagai model.transformer.h[0]. Sisa model (misalnya model.lm_head) dikompilasi secara terpisah.

Hal ini memungkinkan kami untuk mempercepat overhead kompilasi / cold start pada model seperti LLM dan Transformers secara umum. Lihat https://pytorch.org/tutorials/recipes/regional_compilation.html untuk detail lebih lanjut.

Cara Menggunakan Kompilasi Regional

Fitur ini dapat diaktifkan dengan mengatur use_regional_compilation=True dalam konfigurasi TorchDynamoPlugin:

# Configure the compilation backend
dynamo_plugin = TorchDynamoPlugin(
    use_regional_compilation=True,
    ... # other parameters
)
# Initialize accelerator with the plugin
accelerator = Accelerator(dynamo_plugin=dynamo_plugin)
# This will apply compile_regions to your model
model = accelerator.prepare(model)

Anda juga dapat menggunakan utilitas accelerate.utils.compile_regions secara langsung dengan cara yang sama seperti penggunaan torch.compile.

Manfaat Kompilasi Regional

Kami telah melakukan benchmark ekstensif yang membandingkan kompilasi penuh dan kompilasi regional menggunakan fitur torch.compile di PyTorch. Hasil lengkapnya tersedia di repo accelerate. Temuan utama dari benchmark kami adalah:

1. Kinerja yang Setara: Kompilasi regional memberikan peningkatan kecepatan kinerja yang serupa dengan kompilasi penuh, terutama untuk model yang lebih besar.
2. Kompilasi Lebih Cepat: Kompilasi regional secara signifikan mengurangi waktu yang dibutuhkan untuk mengompilasi model, menjadikannya pilihan yang lebih efisien untuk deployment.
3. Dampak Ukuran Batch: Perbedaan kinerja antara strategi kompilasi berkurang seiring dengan ukuran batch yang lebih besar, menunjukkan bahwa overhead kompilasi memiliki dampak yang lebih kecil dalam skenario tersebut.
4. Pertimbangan Ukuran Model: Manfaat kompilasi regional lebih terlihat pada model yang lebih besar, di mana penghematan waktu kompilasi dapat sangat signifikan.
5. Aplikasi Praktis: Untuk aplikasi di dunia nyata, kompilasi regional adalah pilihan praktis untuk mengoptimalkan waktu cold start pelatihan, terutama saat bekerja dengan model besar.

Kesimpulan

Kompilasi penuh maupun regional dapat mempercepat model Anda secara signifikan. Kompilasi regional menawarkan keseimbangan praktis antara waktu kompilasi dan kinerja saat runtime, terutama untuk melatih model besar dengan ukuran batch yang besar.