Internal PyTorch (2019)

• Penjelasan tentang struktur internal PyTorch, sebagai panduan bagi orang-orang yang ingin berkontribusi pada codebase C++ PyTorch
• Tujuan tulisan ini adalah membantu memahami struktur library tensor PyTorch serta teknik diferensiasi otomatis (autograd), dan membantu menavigasi codebase

Struktur dasar tensor PyTorch

• Di PyTorch, tensor adalah struktur data paling dasar
• Tensor adalah struktur data berdimensi-n yang dapat menyimpan nilai skalar seperti floating point (float), integer (int), dan lainnya
• Tensor mencakup metadata berikut:
  • ukuran (size): informasi dimensi tensor
  • dtype: tipe data yang disimpan (misalnya float32, int64, dan lain-lain)
  • device: lokasi tempat data disimpan (CPU, CUDA, dan lain-lain)
  • stride: informasi offset data dalam memori fisik

• Peran stride

  • stride digunakan untuk mengubah indeks logis menjadi lokasi memori fisik
  • stride menetapkan offset untuk setiap dimensi, lalu menentukan lokasi memori fisik dengan mengalikan indeks dengan nilai stride
  • melalui stride, data yang sama dapat dilihat dengan cara berbeda dalam bentuk view tanpa membuat tensor baru

Konsep tensor dan storage

• Di PyTorch, tensor tidak menyimpan data aktual secara langsung → data dikelola di storage
• Tensor = ukuran + dtype + device + stride + offset
• Beberapa tensor dapat berbagi satu storage → mendukung konsep view
• Pemisahan antara storage dan tensor memungkinkan penggunaan memori yang lebih efisien

Proses dispatch operasi tensor

• Di PyTorch, operasi melewati dua tahap dispatch:
  1. Dispatch berdasarkan tipe device dan layout
     • Kode implementasi yang berbeda dijalankan tergantung pada tensor CPU vs. tensor CUDA
  2. Dispatch berdasarkan dtype
     • Kernel yang berbeda dipanggil tergantung pada tipe data seperti float vs. int

Model ekstensi tensor PyTorch

• Tiga elemen ekstensi utama tensor:

  • Device: mendefinisikan cara alokasi memori di CPU, GPU, TPU, dan lainnya
  • Layout: mendefinisikan cara tensor disimpan di memori (misalnya penyimpanan kontigu, penyimpanan sparse, dan lainnya)
  • dtype: mendefinisikan jenis data yang disimpan pada setiap elemen tensor

• Opsi ekstensi:

  • Tensor dapat diperluas dengan memodifikasi langsung kode PyTorch
  • Dapat menulis kelas wrapper yang membungkus tensor yang sudah ada
  • Jika wrapper diperlukan selama diferensiasi otomatis, perlu dilakukan ekstensi langsung

Cara kerja diferensiasi otomatis (Autograd)

• PyTorch melakukan diferensiasi otomatis berdasarkan reverse-mode differentiation
• Saat operasi forward dijalankan, graph dibuat → saat backward, graph ditelusuri untuk melakukan diferensiasi
• Autograd mengelola informasi tambahan berikut:
  • AutogradMeta: metadata yang terhubung ke tensor dan digunakan saat backward
  • Mencatat hasil operasi dan melakukan diferensiasi saat backward

Struktur kode PyTorch dan lokasi file

• Direktori utama dalam codebase PyTorch:
  • torch/ → modul Python (kode Python)
  • torch/csrc/ → kode binding Python dan C++, engine diferensiasi otomatis, compiler JIT, dan lainnya
  • aten/ → definisi operasi tensor (mencakup sebagian besar operasi inti)
  • c10/ → definisi struktur inti seperti tensor dan storage

Proses eksekusi operasi PyTorch

• Contoh: proses eksekusi saat torch.add() dipanggil:
  1. Mengonversi argumen dari Python ke kode C++
  2. Melakukan dispatch di VariableType
  3. Melakukan dispatch berdasarkan device/layout
  4. Menjalankan kernel akhir

Proses penulisan kernel dan alat bantu

• Di PyTorch, kernel ditulis melalui tahapan berikut:
  1. Menulis metadata operasi: mendefinisikan function signature, device yang didukung, dan tipe data
  2. Validasi input: memeriksa dimensi, tipe, dan lainnya
  3. Mengalokasikan tensor output
  4. Dispatch dtype: menjalankan kernel sesuai tipe data
  5. Pemrosesan paralel: di CPU menggunakan OpenMP, di CUDA menggunakan paralelisasi bawaan
  6. Akses dan perhitungan data: menggunakan TensorAccessor, TensorIterator, dan lainnya

Makro dispatch utama

• AT_DISPATCH_ALL_TYPES → melakukan dispatch sesuai dtype
• Dukungan makro tersedia untuk berbagai tipe data → memungkinkan optimasi performa

Optimasi performa dan tips meningkatkan efisiensi kerja

• Minimalkan modifikasi file header → perubahan dapat memicu rebuild seluruh kode
• Siapkan lingkungan pengembangan lokal → meminimalkan waktu yang terbuang jika memakai CI
• Gunakan ccache → dapat menghemat waktu recompilation
• Gunakan server yang kuat → dapat mengurangi waktu kompilasi C++ dan build CUDA

Panduan kontribusi PyTorch

• Target kontribusi yang baik untuk memulai:
  • issue berlabel triaged → issue yang sudah ditinjau oleh pengembang PyTorch
  • membantu perbaikan dokumentasi dan reproduksi bug
  • memberikan masukan terhadap RFC (proposal fitur) PyTorch
• PyTorch telah berkembang melalui kontributor open source, dan partisipasi komunitas sangat disambut