Pembahasan mendalam tentang SparseCore untuk Model Embedding Besar (LEM)

SparseCore adalah prosesor berpetak khusus yang dirancang untuk akselerasi berperforma tinggi pada beban kerja yang melibatkan akses dan komputasi memori sparse yang tidak teratur, terutama pada set data besar yang disimpan dalam High Bandwidth Memory (HBM). Meskipun unggul dalam tugas seperti pencarian sematan, kemampuannya meluas untuk mempercepat berbagai beban kerja dinamis dan jarang lainnya.

1. Pengantar SparseCore

Fitur arsitektur utama:

  • Arsitektur petak: Terdiri dari beberapa petak komputasi (setiap petak adalah unit dataflow lengkap dengan memori lokal dan unit pemrosesannya sendiri) yang memungkinkan pemrosesan paralel.
  • Eksekusi dinamis: Mendukung secara native akses memori dan alur kontrol yang bergantung pada data, yang sangat penting untuk data jarang.
  • Pemrosesan vektor: Memanfaatkan tugas vektor kecil (8 elemen atau 16 elemen, bergantung pada versi hardware) untuk komputasi yang efisien.
  • Kontrol terpusat: Satu pengurut SparseCore mengatur tugas di semua petak, sehingga memastikan operasi yang disinkronkan.
  • Dukungan ringkasan data: Mencakup operasi lintas jalur khusus yang bermanfaat untuk tugas seperti pengurutan, pemfilteran, dan jumlah awalan.
  • Hierarki memori: Memanfaatkan HBM secara strategis untuk menyimpan set data besar dan memori scratchpad lokal (SPMEM) untuk mengatur data yang sering diakses, sehingga mengurangi latensi HBM secara signifikan.

Sekilas tentang spesifikasi:

Atribut TPU v4 TPU v5p Trillium
SparseCores/Chip 4 4 2
Tiles/SparseCore 16 16 16
Lebar SIMD 8 8 8 (F32) 16 (BF16)
Kapasitas HBM 32 GiB 96 GiB 32 GiB

2. Pra-pemrosesan host SparseCore

Persiapan data yang efektif sangat penting untuk performa SparseCore, dan di sinilah praproses host memainkan peran penting. Fitur ini mencakup beberapa fungsi utama:

  • Transformasi data:
    • Terapkan transformasi yang diperlukan pada data input mentah.
    • Mengelola transformasi ID, yang sangat penting saat menangani penumpukan fitur atau tabel.
    • Konversi data input ke dalam format jarang Coordinate (COO), yang dijelaskan di bagian berikut.
    • Partisi data untuk distribusi yang efisien di berbagai SparseCore yang tersedia di chip.
  • Validasi batas:
    • Pastikan karakteristik data input (misalnya, jumlah ID) sesuai dengan batas operasional SparseCore yang telah ditentukan sebelumnya, seperti max_ids_per_partition dan max_unique_ids_per_partition.
    • Jika data input melebihi batas ini, lapisan praproses host dapat mencoba membagi data menjadi mini-batch yang lebih kecil yang sesuai dengan batasan.
  • Transfer data:
    • Menyalin data yang diproses dan divalidasi secara efisien ke High Bandwidth Memory (HBM) TPU, sehingga siap untuk eksekusi SparseCore.

Memahami penumpukan tabel:

Penumpukan tabel adalah teknik pengoptimalan yang signifikan di mana beberapa tabel penyematan digabungkan secara logis untuk meningkatkan efisiensi pencarian penyematan. Proses ini biasanya ditangani secara otomatis oleh framework ML yang mendasarinya.

  • Penumpukan fitur: Hal ini terjadi saat beberapa fitur berbeda berbagi tabel penyematan dasar yang sama. Contoh umumnya adalah menggunakan satu kamus penyematan untuk berbagai fitur kategoris seperti kode pos dari konteks yang berbeda.
  • Penumpukan tabel: Dalam skenario ini, beberapa tabel sematan yang berbeda ditumpuk bersama. Tabel yang memiliki dimensi penyematan dan konfigurasi pengoptimal yang sama sering dikelompokkan.

Keuntungan utama penumpukan tabel adalah pembuatan ukuran batch efektif yang lebih besar untuk operasi pada tabel yang ditumpuk ini. Hal ini mengurangi overhead komputasi dan dapat efektif dalam menyembunyikan latensi komunikasi antar-chip (ICI). Untuk performa yang optimal, sebaiknya gunakan sejumlah tabel bertumpuk yang sedang (biasanya dalam rentang 5 hingga 100).

3. Konversi ke tensor COO

Sebelum data dapat diproses oleh SparseCore, data tersebut biasanya dikonversi ke dalam format tensor jarang (sparse tensor) Coordinate (COO). Format COO adalah cara untuk merepresentasikan matriks jarang secara efisien, biasanya menggunakan tiga array:

  • row_ids: Array yang berisi indeks baris untuk setiap elemen bukan nol. Dalam konteks pemrosesan batch, hal ini sering kali sesuai dengan dimensi batch.
  • col_ids: Array yang berisi indeks kolom untuk setiap elemen non-nol. Untuk penyematan, ini sering kali merupakan nilai fitur atau ID.
  • values (opsional): Array yang menyimpan nilai sebenarnya dari elemen non-nol pada koordinat (row, col) yang sesuai. Untuk perhitungan batas (yang akan dibahas nanti) terkait jumlah ID, nilai ini (keuntungan) sering kali tidak dipertimbangkan.

Contoh ilustrasi:

Pertimbangkan matriks jarang input yang merepresentasikan batch ID:

[
    [id_A],                 // Sample 0
    [id_A, id_B, id_C],     // Sample 1
    [id_B, id_B, id_D],     // Sample 2 (note duplicate id_B)
]

Setelah konversi ke format COO (dan mungkin setelah menghapus duplikat ID dalam sampel yang sama):

row_ids = [0, 1, 1, 1, 2, 2]
col_ids = [id_A, id_A, id_B, id_C, id_B, id_D]

Konversi ini mendasar bagi cara SparseCore memproses dan mendistribusikan pekerjaan. Khususnya, col_ids sangat penting untuk menentukan partisi SparseCore tertentu yang menjadi tujuan ID, sehingga memungkinkan sharding dan pencarian yang efisien.

4. SparsecoreConfig: API tingkat tinggi

API penyematan khusus framework:

SparsecoreConfig, atau mekanisme yang setara seperti flag XLA, berfungsi sebagai antarmuka tingkat tinggi untuk mengontrol berbagai perilaku SparseCore. Pemahaman yang menyeluruh tentang parameter ini sangat penting untuk penyesuaian performa yang efektif dan memastikan pengoperasian model Anda yang benar.

  • disable_table_stacking: bool = False
    • Penjelasan: Flag ini mengontrol apakah penumpukan tabel otomatis mencegah framework menumpuk tabel, yang berpotensi menyebabkan penurunan performa karena peningkatan overhead dan penurunan kemampuan untuk menyembunyikan latensi Inter-Chip Interconnect (ICI).
    • Default: False (menyiratkan penumpukan tabel umumnya diaktifkan secara default jika framework mendukungnya).
  • max_ids_per_chip_per_sample: int = 64
    • Penjelasan: Parameter ini menetapkan batas atas global pada jumlah total ID sematan yang dapat diproses oleh satu chip dari satu sampel dalam batch input, yang digabungkan di semua tabel. Ini adalah mekanisme untuk mengelola resource di tingkat chip, sebelum batas per tabel atau per partisi yang lebih terperinci diperhitungkan. Penyesuaian nilai ini biasanya bergantung pada karakteristik model tertentu dan kapasitas sistem secara keseluruhan.
    • Default: 64.
  • max_ids_per_table: Optional[Dict[str, int]] = None
    • Penjelasan: Parameter ini menentukan jumlah maksimum ID penyematan (yang dapat mencakup duplikat) yang dapat diproses untuk setiap tabel logis, dengan mempertimbangkan semua partisinya di semua SparseCore. Batas ini lebih luas daripada max_ids_per_partition. Jika tabel T dibagi menjadi P partisi, batas ini berlaku untuk jumlah ID yang diarahkan ke semua partisi P. Hal ini sering kali terkait dengan max_ids_per_partition_per_sample dan ukuran batch secara keseluruhan.
    • Setelan: Biasanya dikonfigurasi menggunakan file batas (misalnya, menggunakan tanda xla_sparse_core_max_ids_file), dengan max_ids_per_partition ditentukan. Konsep tingkat tabel ini adalah metode untuk menetapkan batas tingkat partisi tersebut (max_ids dan max_uniques).
    • Default: None (nilai dapat disimpulkan dari batas per partisi atau konfigurasi lain jika tidak diberikan secara eksplisit).
  • max_unique_ids_per_table: Optional[Dict[str, int]] = None
    • Penjelasan: Serupa dengan max_ids_per_table, tetapi parameter ini menentukan jumlah maksimum ID unik untuk setiap tabel logis. Setelan ini sangat penting untuk menentukan ukuran buffer dalam perangkat yang digunakan dalam pemrosesan ID unik dan operasi vektor berikutnya.
    • Setelan: Juga biasanya ditentukan dalam file batas atau berasal dari max_unique_ids_per_partition_per_sample.
    • Default: None.
  • allow_id_dropping: bool = False
    • Penjelasan: Flag boolean ini mengontrol penghapusan ID saat jumlah ID yang ditemukan dalam data input (batas yang diamati) melampaui batas yang ditetapkan selama kompilasi (misalnya, max_ids_per_partition).
      • Jika True: ID yang akan menyebabkan batas terlampaui akan dihapus tanpa pemberitahuan. Biasanya, ID dalam partisi diproses dalam urutan yang diurutkan, dan ID apa pun yang akan mendorong jumlah yang berjalan melebihi batas untuk mini-batch yang ditetapkan akan dihapus. Hal ini memungkinkan program melanjutkan eksekusi, tetapi dapat berdampak buruk pada akurasi model.
      • Jika False: Error akan dipicu, dan proses kemungkinan akan berhenti jika batas yang diamati melampaui batas yang dikompilasi. Pendekatan ini memastikan semua data diproses, tetapi memerlukan konfigurasi batas yang lebih konservatif.
    • Default: False (menyebabkan error saat terjadi overflow, bukan penghapusan data tanpa pemberitahuan).

  • initialize_tables_on_host: bool = True

    • Penjelasan: Flag ini menentukan apakah tabel penyematan diinisialisasi di CPU host sebelum ditransfer ke High Bandwidth Memory (HBM) TPU. Praktik standar adalah tabel diinisialisasi di host. Menetapkan ini ke True mengikuti konvensi ini. Jika disetel ke False, hal ini akan menyiratkan mekanisme inisialisasi di perangkat, yang dapat memiliki implikasi performa yang berbeda atau prasyarat inisialisasi tertentu.

  • enable_fast_table_initialization: bool = False

    • Penjelasan: Menginisialisasi tabel secara langsung di TPU. Hal ini dapat membantu mengurangi waktu mulai model.

5. Pipelining untuk performa

Pipelining adalah teknik pengoptimalan performa yang memungkinkan eksekusi operasi secara bersamaan di TensorCore (TC) dan SparseCore (SC). Dengan tumpang-tindih komputasi ini, throughput secara keseluruhan dapat ditingkatkan secara signifikan.

  • Mekanisme: Dalam langkah pelatihan standar yang melibatkan pencarian sematan jarang (ditangani oleh SC) dan komputasi lapisan padat (ditangani oleh TC), pipelining memungkinkan SC mengerjakan bagian langkah i (misalnya, penerusan atau penerusan mundur) sementara TC secara bersamaan memproses bagian lain dari langkah i yang sama, atau bahkan bagian dari langkah yang berdekatan seperti i-1 atau i+1.
  • Dampak pada gradien: SparseCore mungkin beroperasi pada gradien yang "usang". Misalnya, gradien yang dihitung selama fase backpropagation langkah i mungkin belum sepenuhnya diperbarui dan terlihat oleh SC hingga langkah i+2.
  • Performa vs. trade-off numerik: Eksekusi yang tumpang-tindih ini dapat menghasilkan peningkatan kecepatan yang substansial, yang berpotensi meningkatkan waktu langkah perangkat hingga 2x. Namun, perubahan kecil dalam numerik (embedding_weights) yang dihasilkan dari penggunaan gradien yang tidak berlaku mungkin memengaruhi perilaku konvergensi model atau akurasi akhir yang dicapai. Dapat diterima atau tidaknya pertukaran ini sangat bergantung pada model dan sering kali memerlukan validasi empiris.
  • Flag kontrol: Pipelining dapat dikontrol oleh tf_xla_disable_full_embedding_pipelining. Menyetel flag ini ke true akan menonaktifkan pipelining penuh (komputasi TensorCore dan SparseCore yang tumpang-tindih), sedangkan menyetelnya ke false (atau jika semantik flag menyiratkan pengaktifan saat salah) akan mengaktifkannya.

Alur pipelining konseptual:

  • Tanpa pipelining (alur berurutan yang disederhanakan):

    Loop: SC/F_i -> TC/F_i -> TC/B_i -> SC/B_i

  • Dengan pipelining (alur yang tumpang-tindih yang disederhanakan):

    Time ->
Step i:   SC/F_i | TC/F_i | TC/B_i | SC/B_i
Step i+1:          SC/F_i+1| TC/F_i+1| TC/B_i+1| SC/B_i+1

    Catatan: Tahapan pipelining sebenarnya yang diterapkan di hardware dan compiler bisa lebih rumit, sering kali melibatkan pra-loop, loop eksekusi utama, dan pasca-loop untuk mengelola dependensi data dan memastikan kebenaran.

6. Peran XLA

XLA (Accelerated Linear Algebra) adalah compiler khusus domain yang menerjemahkan grafik komputasi tingkat tinggi, biasanya dari framework seperti TensorFlow, menjadi kode mesin yang sangat dioptimalkan dan disesuaikan untuk TPU. Hal ini mencakup pembuatan petunjuk untuk operasi yang ditujukan untuk SparseCore.

Fungsi utama dalam konteks SparseCore:

  • Kompilasi operasi jarang: XLA bertanggung jawab untuk mengompilasi operasi pencarian penyematan (seperti SparseDenseMatmulOp) dan komputasi jarang lainnya ke dalam program SparseCore tingkat rendah yang dapat dieksekusi.
  • Integrasi batas: Menggunakan batas operasional yang dikonfigurasi (misalnya, max_ids_per_partition, max_unique_ids_per_partition, sering kali diberikan melalui file batas yang ditentukan oleh flag seperti xla_sparse_core_max_ids_file) untuk menentukan ukuran dan mengalokasikan buffer memori di perangkat secara statis, terutama dalam SPMEM.
  • Pengoptimalan yang ditargetkan: XLA melakukan serangkaian pengoptimalan yang dirancang khusus untuk arsitektur SparseCore. Hal ini dapat mencakup penjadwalan instruksi, transformasi tata letak memori, dan penggabungan operasi untuk memaksimalkan efisiensi.
  • Kontrol menggunakan flag: Banyak aspek perilaku SparseCore, parameter penyesuaian, dan strategi pengoptimalan diekspos dan dikontrol melalui flag XLA (misalnya, xla_sparse_core_estimate_max_ids untuk estimasi batas, atau xla_sc_detect_nan untuk proses debug).

Status open source:

Saat ini, Penerapan Sparsecore bersifat internal dan ditayangkan menggunakan libtpu.so.

Pelaporan dan diagnostik error:

Kegagalan kompilasi yang terkait dengan konfigurasi SparseCore atau batasan resource sering kali muncul sebagai XLA:TPUerror waktu kompilasi. Pesan error ini dapat memberikan insight berharga tentang masalah seperti batas yang ditetapkan terlalu tinggi untuk SPMEM yang tersedia, atau penggunaan konfigurasi yang tidak didukung.

7. Cara batas diterjemahkan ke tabel di SparseCore

Di SparseCore, "batas" adalah parameter konfigurasi mendasar yang terutama mengacu pada dua setelan per partisi untuk setiap tabel yang di-shard (didistribusikan) di seluruh SparseCore yang tersedia:

  • max_ids_per_partition: Ini menentukan jumlah maksimum total ID (termasuk duplikat) yang diharapkan dikirim ke, atau diproses oleh, satu SparseCore untuk partisi tertentu dari tabel tertentu dalam satu langkah komputasi.
  • max_unique_ids_per_partition: Ini menentukan jumlah maksimum ID unik yang diharapkan dikirim ke, atau diproses oleh, satu SparseCore

Terjemahan ke tata letak dan pemrosesan tabel fisik:

  • Strategi sharding tabel: Tabel penyematan biasanya "di-shard mod" di semua SparseCore dalam sistem. Artinya, setiap SparseCore bertanggung jawab atas subset kosakata (baris) yang berbeda dari setiap tabel. ID j umumnya akan ditetapkan ke SparseCore_k berdasarkan formula seperti k = j % num_total_sparse_cores.
  • Definisi "partisi": Dalam konteks ini, "partisi" mengacu pada segmen tertentu dari tabel penyematan yang pencariannya ditangani oleh satu SparseCore.
  • Alokasi buffer SPMEM: Batas ini digunakan oleh compiler XLA untuk menentukan ukuran dan mengalokasikan buffer secara statis dalam memori scratchpad di perangkat (SPMEM). Buffer diberi dimensi sehingga semua data yang diperlukan terkait dengan ID untuk partisi tertentu (hingga batas max_ids dan max_unique_ids yang ditentukan) dapat dimuat ke dalam SPMEM untuk diproses. Hal ini sangat penting untuk komputasi non-elementwise, seperti mengurangi ID duplikat dalam partisi (misalnya, , saat membuat representasi Compressed Sparse Row (CSR)), di mana seluruh set data yang relevan untuk ID partisi tersebut harus tersedia dengan cepat di memori cepat.

  • Batas yang dikompilasi versus batas yang diamati:

    • Batas yang diamati: Ini adalah jumlah sebenarnya ID yang ditemukan untuk setiap partisi selama runtime, berdasarkan data input yang diproses.
    • Jika batas yang diamati melebihi batas yang dikompilasi, hal ini dapat menyebabkan ID terlepas (jika allow_id_dropping diaktifkan) atau error.

  • Menghitung batas: Proses penentuan batas yang sesuai melibatkan analisis yang cermat terhadap distribusi data input. Untuk tabel tertentu (sebut saja T1, yang mungkin merupakan bagian dari tabel bertumpuk yang lebih besar T):

    1. Batch input (misalnya, SparseTensor 2D dengan bentuk [BatchSize, MaxSequenceLength]) awalnya dibagi di seluruh SparseCore yang tersedia. Misalnya, jika TensorCore dipasangkan dengan 2 SparseCore, setiap SparseCore dapat menerima sub-batch dengan bentuk [BatchSize/2, MaxSequenceLength].
    2. Sub-batch ini kemudian dikonversi menjadi format COO, sehingga menghasilkan row_ids dan col_ids.
    3. ID duplikat dalam sampel yang sama (yaitu, entri dengan row_id dan col_id yang sama) akan dihapus.
    4. Untuk setiap col_id unik yang tersisa (dalam sampel), SparseCore target yang bertanggung jawab atas ID ini ditentukan menggunakan aturan mod-sharding: target_sc_id = col_id % num_total_sparse_cores.
    5. Jumlah total ID (ids_per_sparse_core[target_sc_id]++) dan jumlah ID unik (unique_ids_per_sparse_core[target_sc_id]++, setelah memastikan keunikan untuk target_sc_id tertentu tersebut) yang ditujukan untuk setiap target_sc_id akan dipertahankan.
    6. max_ids_per_partition untuk tabel T1 kemudian ditetapkan ke max(ids_per_sparse_core_array).
    7. Demikian pula, max_unique_ids_per_partition untuk tabel T1 disetel ke max(unique_ids_per_sparse_core_array).
    8. Jika tabel T1 adalah komponen tabel bertumpuk, transformasi tambahan seperti rotasi atau pergeseran mungkin diterapkan pada distribusi ID sebelum menjumlahkan statistik dari semua tabel penyusun. Hal ini membantu menyeimbangkan beban di seluruh chip.

Menyetel batas ini dengan benar adalah tindakan menyeimbangkan: batas yang lebih rendah berpotensi menghasilkan performa yang lebih tinggi (karena lebih sedikit data yang perlu diproses per langkah dan tekanan SPMEM berkurang), tetapi jika disetel terlalu rendah, batas tersebut dapat menyebabkan mini-batching yang berlebihan atau penghapusan ID yang tidak diinginkan.

8. Cara setiap SparseCore berkomunikasi

Komunikasi SparseCore, khususnya dalam konteks pemrosesan daftar ID untuk pencarian sematan, mengandalkan beberapa mekanisme terkoordinasi:

  • Mod sharding dan pemilihan rute implisit:
    • Tabel sematan di-shard mod di semua SparseCore dalam sistem.
    • Saat host menyediakan batch data input (yang selanjutnya diproses menjadi format COO, termasuk col_ids), nilai col_id digunakan untuk menentukan SparseCore mana yang bertanggung jawab atas ID tertentu tersebut: target_sc_id = col_id % num_total_sparse_cores.
    • Setiap SparseCore secara efektif menerima dan memproses hanya subset ID yang dipetakan ke partisi kosakata yang ditetapkan. Tahap pra-pemrosesan host sangat penting untuk menyiapkan data sedemikian rupa sehingga setiap SparseCore dapat dengan mudah mengidentifikasi dan mengoperasikan ID yang relevan.
  • Distribusi data menurut host:
    • Logika praproses host mempartisi batch input keseluruhan dan mendistribusikan bagian row_ids dan col_ids yang relevan (bersama dengan fitur atau bobot terkait jika berlaku) baik ke memori (HBM) yang dapat diakses langsung oleh setiap SparseCore atau ke HBM bersama yang akan digunakan SparseCore untuk mengambil data yang diperlukan.
  • Pemrosesan Intra-SparseCore:
    • Setelah menerima kumpulan ID yang ditetapkan untuk partisi tabel tertentu, SparseCore akan melakukan operasi seperti penghapusan duplikat ID ini dan pengumpulan vektor sematan yang sesuai. Ini terutama adalah komputasi lokal yang dieksekusi dalam petak SparseCore sendiri dan memanfaatkan SPMEM lokalnya.
  • Komunikasi antar-SparseCore (Semua-ke-Semua):
    • Setelah fase pemrosesan awal (seperti pencarian sematan), pola komunikasi "semua-ke-semua" dapat digunakan untuk menggabungkan atau mendistribusikan ulang hasil di seluruh SparseCore (misalnya, sebelum memasukkan aktivasi ke dalam lapisan TensorCore yang mengharapkan input yang sesuai dengan semua posisi sampel asli). Hal ini sangat penting untuk merekonstruksi kumpulan aktivasi lengkap jika batch input asli didistribusikan untuk pemrosesan paralel.
  • Komunikasi dengan TensorCore:
    • SparseCores berkomunikasi dengan TensorCores untuk mengirim aktivasi penyematan (selama forward pass) dan untuk menerima gradien (selama backward pass). Interaksi ini diatur oleh program yang dikompilasi XLA dan sering kali melibatkan HBM sebagai buffer perantara. Strategi pipelining (yang dibahas sebelumnya) sangat memengaruhi pengaturan waktu dan sinkronisasi komunikasi SC-TC ini.

Pada dasarnya, "distribusi" awal ID ke SparseCore yang sesuai sebagian besar ditangani oleh skema sharding dan langkah-langkah praproses host. Komunikasi berikutnya melibatkan SparseCore yang beroperasi pada data lokalnya, yang berpotensi diikuti oleh operasi komunikasi kolektif seperti all-to-all jika data perlu dipertukarkan atau diurutkan ulang secara global di seluruh SparseCore sebelum diproses lebih lanjut oleh TensorCore.

9. Pengelolaan memori SparseCore

Setiap SparseCore secara efisien mengelola beberapa jenis memori yang berbeda untuk melakukan komputasinya:

  • Memori scratchpad (SPMEM):
    • Alam: SRAM lokal yang relatif kecil, tetapi sangat cepat, yang tersedia secara eksklusif untuk setiap SparseCore. Penting untuk diperhatikan bahwa SPMEM bukan cache; penggunaannya dikelola dan diatur secara eksplisit oleh compiler XLA.
    • Tujuan: SPMEM digunakan untuk "menyiapkan data secara oportunistik". Hal ini mencakup input, output, dan hasil perantara yang diperlukan untuk komputasi SC yang sedang berlangsung. Menyimpan data dalam SPMEM secara signifikan mengurangi latensi tinggi yang biasanya terkait dengan akses HBM.
    • Penentuan ukuran: Seperti yang dibahas di bagian "Batas", buffer SPMEM ditentukan ukurannya secara statis pada waktu kompilasi. Penentuan ukuran ini didasarkan pada parameter seperti max_ids_per_partition dan max_unique_ids_per_partition. Alokasi statis ini memastikan bahwa untuk setiap operasi tertentu pada partisi tabel (seperti pengurangan CSR), semua data yang diperlukan untuk ID partisi tersebut (hingga batas yang ditentukan) dapat masuk ke SPMEM.
    • Pengoptimalan compiler: Compiler XLA menggabungkan pengoptimalan canggih untuk menentukan secara tepat berapa banyak data, dan elemen data spesifik mana, yang perlu diatur dalam SPMEM untuk menyembunyikan latensi HBM secara efektif dan memaksimalkan performa.
    • Batasan alokasi dinamis: Compiler SparseCore saat ini tidak mendukung alokasi ruang kerja sementara dinamis. Hal ini menyoroti pentingnya ukuran statis melalui konfigurasi batas yang cermat.
  • High bandwidth memory (HBM):
    • Nature: Resource memori bersama yang besar dan dapat diakses oleh semua SparseCore, TensorCore, dan sistem host. Tabel sematan utama disimpan di HBM.
    • Penggunaan stack: Operasi SparseCore sering kali memerlukan penyimpanan sementara di HBM untuk hasil perantara yang tidak sesuai dengan SPMEM terbatas atau perlu diteruskan di antara tahap yang lebih besar dalam pipeline pemrosesan. Penggunaan stack HBM selama lintasan maju dan mundur dapat diperkirakan sebagai berikut -
      • Stack HBM Forward Pass (tabel tunggal) ≈ (2 * feature_width + 1) * max_unique_nz_per_row * logical_replica_count * 4 byte
      • Stack HBM Backward Pass (tabel tunggal) ≈ 3 * feature_width * max_unique_nz_per_row * logical_replica_count * 4 byte
    • Penggunaan heap: HBM juga mengakomodasi heap, yang dikelola oleh host. Heap menyimpan data seperti bobot lapisan padat, konstanta yang digunakan oleh model, dan data input yang telah diambil sebelumnya. Penggunaan heap cenderung meningkat seiring dengan jumlah langkah yang digunakan host untuk melakukan pengambilan data terlebih dahulu (dikontrol oleh tanda maximum_parallel_iterations). Meskipun pengambilan data lebih awal dapat meningkatkan performa dengan tumpang-tindih transfer host-ke-perangkat dengan komputasi perangkat, hal ini juga menggunakan lebih banyak HBM.
    • Serialisasi untuk pengoptimalan HBM: Flag xla_sc_num_serialized_tables_to_optimize_hbm menyediakan mekanisme untuk mengontrol jumlah data tabel yang tetap "aktif" dalam memori stack HBM pada waktu tertentu. Meningkatkan jumlah ini secara efektif akan menserialkan pemrosesan untuk lebih banyak tabel, yang dapat mengurangi penggunaan stack HBM puncak, tetapi dapat mengurangi performa karena berkurangnya paralelisme.
  • Memori vektor (VMEM):
    • VMEM adalah memori scratchpad lokal yang digunakan secara eksklusif oleh TC (TensorCore). Meskipun tidak dikelola langsung oleh SparseCore, VMEM adalah bagian integral dari ekosistem memori yang berinteraksi dengan SC, terutama melalui TensorCore.

Strategi pengelolaan memori secara keseluruhan:

Strategi pengelolaan memori inti untuk SparseCore berkisar pada penggunaan SPMEM kecil dan cepat untuk data "aktif" yang sedang diproses oleh petak SparseCore, sehingga meminimalkan akses ke HBM yang lebih lambat. Batas yang dikonfigurasi adalah mekanisme utama untuk memastikan bahwa SPMEM tidak meluap. HBM digunakan untuk menyimpan tabel embedding besar dan data sementara yang melebihi kapasitas SPMEM atau perlu dibagikan di berbagai unit pemrosesan atau tahap pipeline. Compiler XLA bertanggung jawab untuk mengatur semua pergerakan data dan alokasi buffer berdasarkan prinsip arsitektur ini dan batas yang dikonfigurasi pengguna.

10. Bottleneck performa & memori

Untuk mencapai performa optimal dengan SparseCore, Anda harus memahami dengan jelas potensi hambatan dan cara mengatasinya. Masalah ini dapat muncul di host, dalam SparseCore itu sendiri, atau dalam interaksinya dengan TensorCore.

Hambatan performa umum:

  • Hambatan host:
    • Masalah: CPU host mungkin gagal memproses data sebelumnya dan mengirimkannya ke TPU dengan cukup cepat, sehingga menyebabkan kurangnya pemanfaatan SparseCore dan TensorCore. Hal ini sering kali membatasi performa.
    • Mitigasi: Pantau pemakaian CPU host dan metrik pipeline input. Mengoptimalkan rutin pemuatan dan praproses data sisi host (lihat tips konversi COO). Sesuaikan tanda maximum_parallel_iterations untuk menyempurnakan pengambilan data sebelumnya.
  • Sinkronisasi TC/SC yang tidak optimal (kurangnya pipelining):
    • Masalah: Jika pipelining antara TensorCore dan SparseCore dinonaktifkan atau tidak beroperasi secara efisien, satu unit mungkin menghabiskan banyak waktu untuk menunggu unit lainnya, sehingga mengurangi throughput sistem secara keseluruhan.
    • Mitigasi: Pastikan pipelining diaktifkan (misalnya, tf_xla_disable_full_embedding_pipelining = false atau yang setara).
  • Hambatan yang disebabkan oleh batas:
    • Masalah:
      • Batas terlalu rendah: Dapat memicu mini-batching yang berlebihan (membagi batch input menjadi banyak sub-batch yang lebih kecil untuk memenuhi batas yang ketat). Meskipun hal ini mempertahankan kebenaran, setiap mini-batch memperkenalkan beberapa overhead pemrosesan, yang berpotensi memperlambat eksekusi secara keseluruhan. Jika allow_id_dropping benar, batas yang terlalu rendah juga dapat menyebabkan ID dihapus, yang memengaruhi akurasi model.
      • Batas terlalu tinggi (tetapi masih sesuai): Meskipun batas yang sangat tinggi dapat mencegah pengelompokan mini-batch, batas tersebut dapat meningkatkan tekanan SPMEM secara tidak perlu jika karakteristik data sebenarnya jarang mendekati nilai puncak ini. Hal ini juga dapat menyebabkan penggunaan stack HBM yang lebih besar daripada yang diperlukan.
      • Kegagalan kompilasi: Jika batas yang dikonfigurasi memerlukan stack SPMEM atau HBM yang lebih banyak daripada memori fisik yang tersedia, kompilasi akan gagal.
    • Mitigasi: Pastikan batas ditetapkan dengan benar.
  • Ketidakseimbangan distribusi data:
    • Masalah: Jika partisi SparseCore tertentu secara konsisten menerima jumlah ID yang lebih besar secara tidak proporsional dibandingkan dengan yang lain (menunjukkan distribusi ID yang buruk), SparseCore yang kelebihan beban tersebut akan menjadi hambatan performa.
    • Mitigasi: Pengacakan ID selama proses mini-batching dapat membantu mengurangi masalah ini untuk tabel bertumpuk, terutama yang memiliki tabel pengguna "aktif". Analisis distribusi ID dengan cermat untuk menetapkan batas per tabel yang sesuai dan seimbang.
  • Masalah penumpukan tabel:
    • Masalah:
      • Terlalu sedikit tabel yang ditumpuk: Mungkin tidak cukup untuk menyembunyikan latensi ICI secara efektif atau mengurangi overhead pemrosesan secara memadai.
      • Terlalu banyak tabel yang ditumpuk: Dapat mengakibatkan pembuatan tabel logis yang sangat besar sehingga sulit dikelola atau dapat melebihi batas resource yang tersedia.
    • Mitigasi:
      • Pastikan jumlah tabel yang optimal untuk penumpukan. Panduan umum menyarankan "titik ideal" 5-100 tabel untuk penumpukan.
  • Numerik/kuantisasi yang tidak efisien:
    • Masalah: Menggunakan presisi FP32 penuh saat format presisi yang lebih rendah seperti BF16 atau bilangan bulat terkuantisasi sudah cukup (dan menawarkan komputasi yang lebih cepat) dapat menjadi hambatan performa.
    • Mitigasi: Jelajahi opsi presisi yang lebih rendah. Namun, perlu diketahui bahwa kuantisasi itu sendiri memiliki beberapa overhead dan mungkin memerlukan penyesuaian parameter kuantisasi yang cermat untuk mempertahankan akurasi model.
  • Saturasi bandwidth HBM:
    • Masalah: Pergerakan data yang berlebihan ke dan dari HBM, yang berpotensi disebabkan oleh lebar fitur yang sangat kecil (sehingga menyebabkan overhead padding yang tinggi), pola akses memori yang tidak efisien, atau jumlah pencarian yang sangat besar, dapat membuat bandwidth HBM yang tersedia menjadi jenuh.
    • Mitigasi: Menskalakan jumlah TPU dapat membantu mengatasi saturasi bandwidth HBM.

Hambatan memori umum:

  • SPMEM overflow (kegagalan kompilasi):
    • Masalah: Jika max_ids_per_partition dan max_unique_ids_per_partition disetel terlalu tinggi, compiler XLA mungkin tidak dapat mengalokasikan SPMEM yang cukup, sehingga menyebabkan error kompilasi seperti: "Fixed size allocations (...) do not fit in TileSpmem (...)". Selain itu, jika istilah (sample_count * feature_width) / kNumTiles (dengan kNumTiles adalah jumlah petak per SC) terlalu besar untuk mengumpulkan operan penyiapan dalam SPMEM petak, error seperti "Gather operand too large..." dapat terjadi.
    • Mitigasi: Kurangi ukuran batch atau tingkatkan jumlah chip yang digunakan untuk pemrosesan.
  • Stack overflow HBM (runtime atau kompilasi):
    • Masalah: Jika kombinasi feature_width, max_unique_nz_per_row, dan logical_replica_count menyebabkan persyaratan memori stack HBM yang melebihi HBM yang tersedia, hal ini dapat menyebabkan error Kehabisan Memori (OOM) baik saat runtime maupun selama kompilasi.
    • Mitigasi: Sesuaikan tanda xla_sc_num_serialized_tables_to_optimize_hbm untuk mengurangi penggunaan stack HBM dengan melakukan serialisasi pemrosesan tabel (biasanya akan memengaruhi performa).
  • Penggunaan heap HBM yang berlebihan:
    • Masalah: Terutama disebabkan oleh bobot lapisan padat yang sangat besar, banyak konstanta yang disimpan dalam memori, atau pengambilan data input yang terlalu agresif (maximum_parallel_iterations tinggi).
    • Mitigasi: Pantau penggunaan heap menggunakan alat seperti XProf Memory Viewer.
  • Overhead padding:
    • Masalah: Tabel sematan diisi agar selaras dengan 32B (setara dengan 8 float) dalam dimensi fitur. Akibatnya, lebar fitur kecil (misalnya, 1 float) menimbulkan overhead padding yang signifikan (misalnya, 7/8 ruang buffer yang dialokasikan adalah padding), sehingga menyebabkan HBM terbuang. Dimensi kosakata tabel juga diisi agar menjadi kelipatan jumlah SparseCore dalam sistem; namun, dampak ini biasanya dapat diabaikan untuk tabel dengan ukuran kosakata yang cukup tinggi.

Faktor umum yang memengaruhi performa dan memori:

  • Topologi: Jumlah chip yang tersedia dan arsitektur interkoneksinya.
  • Ukuran batch: Memengaruhi sample_count per SparseCore secara langsung, yang pada gilirannya memengaruhi konsumsi memori dan beban komputasi.
  • Pemformatan data: Memastikan tata letak data di perangkat yang efisien sangat penting untuk performa yang optimal.

11. Menganalisis profil SparseCore

Menganalisis profil performa adalah langkah penting dalam mengidentifikasi hambatan dan menemukan peluang pengoptimalan dalam workload SparseCore Anda.

  1. Mendapatkan rekaman aktivitas:
    • Gunakan alat pembuatan profil, seperti XProf, untuk merekam aktivitas eksekusi yang mendetail saat model Anda dilatih atau menjalankan inferensi. Rekaman aktivitas ini akan memberikan linimasa operasi yang terjadi di host, TensorCore, dan SparseCore.
  2. Periksa Trace Viewer (misalnya, di XProf atau TensorBoard):
    • Aktivitas host: Memeriksa aktivitas host. Apakah ada kesenjangan yang signifikan dalam aktivitas TPU? Kesenjangan tersebut mungkin menunjukkan bahwa host menjadi hambatan, karena gagal mengirimkan data dengan cukup cepat. Analisis performa pipeline input Anda.
    • Aktivitas TensorCore (TC) dan SparseCore (SC):
      • Lihat linimasa eksekusi untuk TC dan SC. Apakah keduanya beroperasi secara paralel, yang menunjukkan pipelining yang efektif? Atau, apakah ada periode yang lebih lama saat satu unit tidak aktif, menunggu unit lainnya?
      • Identifikasi operasi yang menghabiskan waktu paling banyak (operasi yang berjalan paling lama) di SC dan TC.
      • Output rekaman visual (sering kali menampilkan blok berwarna yang merepresentasikan berbagai operasi dari waktu ke waktu, seperti TPU:0 SparseCore 1 (pid 1005)) sangat berharga untuk mengidentifikasi secara visual operasi dominan dan periode tidak ada aktivitas.
    • Analisis waktu langkah: Amati waktu langkah secara keseluruhan dan pahami cara pendistribusian atau perinciannya antara pemrosesan host, komputasi SC, dan komputasi TC.
  3. Analisis memori (XProf Memory Viewer):
    • Penggunaan heap: Gunakan alat seperti tab "Memory Viewer" XProf untuk memeriksa penggunaan heap HBM. Hal ini dapat membantu menentukan apakah bobot model besar, konstanta, atau pengambilan data input yang terlalu agresif menggunakan HBM dalam jumlah yang berlebihan. Mengaktifkan flag seperti --vmodule=best_fit_allocator=1 dapat memberikan log penggunaan heap puncak.
    • Penggunaan stack (tidak langsung): Meskipun pembuatan profil stack HBM langsung bisa jadi rumit, jika Anda mengalami error Kehabisan Memori dan penggunaan heap tampak wajar, kehabisan stack HBM (sering kali karena batas atau lebar fitur yang terlalu besar) adalah penyebab yang paling mungkin. Formula yang disediakan untuk penggunaan stack HBM dapat membantu memperkirakan hal ini.
  4. Cari pola tertentu:
    • Tumpukan mini: Jika batas sering terlampaui, Anda mungkin mengamati bukti tumpukan mini dalam rekaman aktivitas (misalnya, jumlah operasi SC yang lebih kecil lebih banyak dari yang diharapkan untuk ukuran tumpukan global). Hal ini sering kali dapat disimpulkan dari log atau dengan mengamati jumlah pemanggilan operasi tertentu.
    • Pelepasan ID: Jika pelepasan ID diaktifkan dan terjadi, log sistem mungkin memberikan indikasi hal ini. Hal ini juga akan menjadi tanda yang jelas bahwa batas yang dikonfigurasi terlalu ketat untuk data input.
    • Waktu kompilasi: Waktu kompilasi yang diperpanjang, terutama jika Pengoptimalan yang Diarahkan Masukan (FDO) diaktifkan dan sering menyesuaikan batas, dapat menambah overhead yang signifikan pada waktu pelatihan secara keseluruhan.
  5. Korelasi dengan tanda dan konfigurasi:
    • Hubungkan perilaku yang diamati dalam profil kembali ke konfigurasi SparseCore Anda (setelan dalam file batas, flag XLA). Misalnya, jika xla_sc_num_serialized_tables_to_optimize_hbm ditetapkan ke nilai yang tinggi, Anda mungkin mengharapkan performa SC yang lebih lambat, tetapi penggunaan stack HBM yang lebih rendah.
  6. Proses iteratif:
    • Pembuatan profil sering kali merupakan proses penyempurnaan iteratif. Lakukan perubahan tertentu (sesuaikan batas, aktifkan atau nonaktifkan fitur), ambil profil baru, lalu bandingkan dengan profil sebelumnya untuk melihat dampak modifikasi Anda.

12. Flag proses debug umum

Beberapa tanda dapat diaktifkan untuk membantu proses debug masalah terkait eksekusi SparseCore. Penting untuk diperhatikan bahwa mengaktifkan pemeriksaan ini sering kali menimbulkan penalti performa dan, oleh karena itu, pemeriksaan ini biasanya harus dinonaktifkan untuk menjalankan produksi.

  • Pemeriksaan ID (di luar rentang):
    • Bendera: xla_sparse_core_enable_id_bound_check = true
    • Tujuan: Memungkinkan pemeriksaan pada sistem host untuk mendeteksi apakah ada ID penyematan dalam data input yang berada di luar rentang kosakata yang valid yang ditentukan untuk tabel penyematan tertentu. Hal ini membantu mendeteksi masalah terkait data input yang salah atau rusak.
  • Pemeriksa NaN:
    • Bendera: xla_sc_detect_nan = true
    • Tujuan: Memungkinkan deteksi nilai NaN (Not a Number) dalam data floating point yang diproses di SparseCore. Jika NaN terdeteksi dalam input atau output berbagai langkah kompilator, tanda ini akan menyebabkan error muncul. Error tersebut biasanya memberikan informasi tentang tempat NaN ditemukan.
  • Pemeriksa batas (akses memori):
    • Bendera: xla_sc_assert_level=bounds
    • Tujuan: Flag ini mengaktifkan alat gaya ASAN (AddressSanitizer) yang menulis ulang instruksi akses memori (seperti operasi pemuatan/penyimpanan VMEM dan DMA) untuk menyertakan pemeriksaan dinamis. Pemeriksaan ini memverifikasi apakah akses memori berada dalam batas yang dialokasikan dari wilayah memori target.
    • Perilaku: Jika akses memori di luar batas terdeteksi, eksekusi akan gagal.
    • Perhatian: Pemeriksa ini dapat menghasilkan positif palsu, misalnya, karena pola akses ber-stride yang kompleks yang tidak sepenuhnya dipahami oleh pemeriksa. Transformasi ini diterapkan pada tahap akhir dalam proses kompilasi backend.
  • Pemeriksa buffer (kerusakan memori):
    • Flag:
      • xla_tpu_buffer_contents_sanitizer_config='cores_to_sanitize: [TC, SC_SCS, SC_TILE], sanitizer_mode: LOCAL_ONLY'
      • xla_tpu_verify_launch_id_across_cores=true
    • Tujuan: Flag ini membantu memastikan bahwa buffer memori tidak rusak atau ditimpa secara tidak sengaja oleh operasi yang tidak terkait. Buffer Sanitizer memeriksa konten buffer untuk memverifikasi bahwa konten tersebut tidak berubah secara tidak terduga.

13. Dukungan kuantisasi

SparseDenseMatmulOp SparseCore dirancang untuk mendukung operasi pada tabel penyematan menggunakan jenis data floating-point 32-bit (FP32) dan bilangan bulat. Meskipun pelatihan model biasanya dilakukan menggunakan presisi FP32 untuk tabel penyematan, kuantisasi pasca-pelatihan (PTQ) dapat diterapkan. PTQ memungkinkan penggunaan jenis data presisi yang lebih rendah (seperti bilangan bulat 8-bit) untuk inferensi, yang berpotensi meningkatkan performa dan mengurangi footprint memori.

Simulasi Kuantisasi:

SparseDenseMatmulOp dapat dikonfigurasi untuk melakukan "kuantisasi simulasi". Dalam mode operasional ini, vektor embedding pertama-tama dikuantisasi ke presisi yang lebih rendah, lalu didekuantisasi kembali ke presisi yang lebih tinggi (misalnya, FP32) sebelum digunakan dalam komputasi berikutnya. Teknik ini memungkinkan model dilatih sambil memperhitungkan efek derau kuantisasi. Pelatihan dengan kuantisasi simulasi dapat meningkatkan akurasi model akhir saat dikuantisasi sepenuhnya untuk inferensi.

Atribut konfigurasi untuk SparseDenseMatmulOp (untuk kuantisasi):

  • quantization_config_num_buckets = 256
    • Atribut ini menentukan jumlah bucket atau tingkat diskrit yang akan digunakan untuk menguantisasi bilangan floating point 32-bit. Misalnya, saat melakukan kuantisasi ke bilangan bulat 8-bit, biasanya akan ditentukan 2^8 =256 bucket.
  • quantization_config_low = -X.X
    • Atribut ini menentukan nilai floating point minimum dalam rentang kuantisasi. Setiap nilai input di bawah minimum yang ditentukan ini akan dipangkas ke nilai minimum ini selama kuantisasi.
  • quantization_config_high = Y.Y
    • Atribut ini menentukan nilai floating point maksimum dalam rentang kuantisasi. Nilai input apa pun di atas maksimum yang ditentukan ini akan dipangkas ke nilai maksimum ini selama kuantisasi.

Interaksi numerik dan pipelining:

Perilaku numerik model dapat berubah, bergantung pada apakah pipelining antara TensorCore dan SparseCore diaktifkan. Jika pipelining aktif, gradien yang diproses oleh SparseCore mungkin "usang" (dari iterasi sebelumnya). Hal ini dapat berinteraksi dengan proses kuantisasi dan berpotensi memengaruhi dinamika pelatihan model atau akurasi akhir.

14. Fitur mendatang dan peningkatan terkini

Ekosistem SparseCore tunduk pada pengembangan dan peningkatan berkelanjutan.

Peta jalan:

  • Pengelompokan mini menurut dimensi sampel:
    • Fitur ini direncanakan sebagai fitur pelengkap untuk kemampuan batch mini dimensi kosakata yang ada.
    • Hal ini akan memungkinkan partisi lebih lanjut dari input embedding di sepanjang dimensi sampel. Hal ini akan dicapai dengan memperkenalkan loop di perangkat yang dapat memfilter dan memproses pencarian dari subset sampel dalam satu waktu. Fitur ini dapat bermanfaat untuk mengelola jumlah ID per sampel yang sangat besar atau untuk meningkatkan penyeimbangan beban di seluruh unit pemrosesan.
  • Peningkatan dukungan untuk penyematan dengan kurang dari 8 bilangan bulat per baris (lebar fitur kecil):
    • Desain saat ini sering menggunakan padding yang signifikan untuk menyematkan lebar fitur yang kurang dari 8 float (yang sesuai dengan 32 byte). Pengisihan ini dapat menyebabkan HBM terbuang dan berpotensi kurangnya pemanfaatan resource komputasi. Peningkatan di masa mendatang bertujuan untuk mengurangi inefisiensi ini untuk tabel dengan dimensi fitur kecil.

Peningkatan terbaru:

  • Penyiapan operand pengumpulan di HBM:
    • Pengoptimalan ini membantu mengurangi tekanan pada memori scratchpad bersama (SPMEM) dengan mengizinkan beberapa input atau output operasi pengumpulan untuk diatur dalam HBM yang lebih besar.
  • Pengurangan penggunaan memori stack:
    • Peningkatan telah diterapkan untuk mengurangi konsumsi memori stack HBM selama operasi SparseCore, idealnya tanpa berdampak negatif pada performa atau throughput secara keseluruhan.

Peningkatan ini berfokus pada peningkatan performa, efisiensi memori, dan fleksibilitas operasional SparseCore untuk berbagai beban kerja jarang yang lebih luas.