Model Biaya LHS

tldr;

Halaman ini menjelaskan internal model biaya yang digunakan oleh Latency Hiding Scheduler. Jika Anda tertarik untuk menyetel model, langsung buka bagian Penyetelan.

Latency Hiding Scheduler (LHS) adalah compiler pass yang menjadwalkan DAG HLO sedemikian rupa sehingga meminimalkan waktu nyata.

Keputusannya dipandu oleh model biaya terpadu, yang menggunakan campuran tabel performa dan model analitis. Secara khusus, XLA menyematkan tabel performa untuk GEMM dan kolektif interkoneksi cepat, serta menggunakan model biaya fusi dan jaringan analitik untuk kasus lainnya. Bagian lain dokumen ini menjelaskan cara kerja internalnya secara umum.

Tabel performa – koleksi ICI

Tabel performa terdiri dari dua komponen utama: pengumpul dan interpolator.

Kolektor

Pengumpul (collector) adalah alat C++ yang bertanggung jawab untuk membuat tabel performa untuk operasi kolektif. Mengukur performa setiap operasi HLO (misalnya, all-gather, all-reduce) di seluruh ruang parameter yang ditentukan secara statis.

Cara Kerjanya

Alat ini melakukan pemindaian pada berbagai operasi kolektif, ukuran transfer, dan skema transfer untuk cluster tertentu. Infrastruktur ini menggunakan runner HLO multi-host yang ada dan data ExecutionProfile untuk menjalankan HLO yang dihasilkan dan mengumpulkan metrik performa.

Parameter Pengumpulan Data

Tabel latensi dikumpulkan untuk produk silang dari parameter berikut:

  • Collective Type:
    • all-reduce
    • all-gather
    • reduce-scatter
  • Ukuran Transfer:
    • Skala logaritmik dari 1024 B hingga 2 GiB (misalnya, 1024B, 2048B, 4096B, ...)
  • Skema Transfer:
    • rail-aligned
    • non-rail-aligned

Penyapuan ini dijalankan untuk cluster intra-node dengan 2, 4, dan 8 perangkat.

Output

Hasil dari eksekusi pengumpulan adalah tabel latensi dalam format .pbtxt (sekitar 116 KB per platform).

Interpolator

Interpolator adalah komponen compiler yang menggunakan tabel performa yang dihasilkan untuk memberikan perkiraan runtime selama kompilasi.

Struktur Data Internal

Saat inisialisasi, Interpolator memproses tabel performa menjadi peta. Peta ini menggunakan tuple (collective_type, transfer_scheme) sebagai kuncinya.

Nilai yang terkait dengan setiap kunci adalah bidang Euklidean 2D. Bidang ini mengindeks throughput jaringan (diukur oleh Pengumpul) berdasarkan dua sumbu:

  1. Ukuran transfer.
  2. Jumlah perangkat yang terlibat.

Pencarian dan Interpolasi

Saat compiler menemukan operasi kolektif, Interpolator akan melakukan langkah-langkah berikut:

  1. Fungsi ini mengidentifikasi bidang throughput 2D yang benar menggunakan (collective_type, transfer_scheme) operasi sebagai kunci peta.
  2. Kemudian, menggunakan pengambilan rata-rata berbobot (berdasarkan jarak Euclidean) dalam bidang 2D tersebut, menggunakan (transfer_size, num_devices) operasi sebagai titik kueri.
  3. Hasil pencarian ini adalah satu nilai throughput jaringan yang unik.

Alasan: Throughput dan Ekstrapolasi

Sistem ini dirancang untuk menyimpan throughput jaringan, bukan latensi mentah. Pilihan desain ini secara signifikan menyederhanakan ekstrapolasi performa untuk ukuran transfer yang tidak ada secara eksplisit dalam tabel.

Jika tabel latensi merekam saturasi bandwidth jaringan pada ukuran kolektif S, throughput T pada saat itu dianggap sebagai maksimum. Untuk setiap koleksi baru berukuran S' > S, waktu proses dapat diperkirakan sebagai:

\[\text{EstimatedTime}(S') = \frac{S'}{T_{\text{saturated} } }\]

Hal ini memungkinkan model memperkirakan performa untuk kolektifitas dengan ukuran apa pun, bahkan yang lebih besar dari maksimum 2 GiB yang diukur oleh Pengumpul.

  • Meremehkan throughput maksimum.
  • Akibatnya, perkirakan berlebih runtime untuk transfer besar.

Secara umum, tim XLA:GPU mengelola tabel performa, tetapi jika pengguna memutuskan untuk menyediakan tabelnya sendiri, pengguna yang membuat tabel tersebut bertanggung jawab untuk memastikan tabel tersebut representatif dan menyertakan pengukuran di wilayah yang jenuh bandwidth untuk hardware target.

Tabel performa – GEMM

Mirip dengan sistem untuk kolektif, tabel latensi GEMM didukung oleh dua komponen: pengumpul dan penginterpolasi.

Kolektor

Pengumpul (collector) adalah alat C++ yang menghitung tabel performa untuk Perkalian Matriks Umum (GEMM). Benchmark ini mengukur performa perkalian matriks di tingkat operasi HLO dot.

Cara Kerjanya

Alat ini melakukan pemindaian di ruang statis dimensi GEMM (batch, dua dimensi non-kontraksi, dan satu dimensi kontraksi) serta jenis data.

  • Jenis Data Default: LHS = bf16,f32, RHS = bf16,f32, OUT = bf16,f32.
  • Infrastruktur: Menggunakan kembali profiler op HLO.

Parameter Pengumpulan

Tabel latensi dikumpulkan untuk produk gabungan dari dimensi berikut:

  • batch: {1, 2, 4}
  • m (non-kontrak): {256, 512, ..., 4096}
  • n (non-kontrak): {256, 512, ..., 4096}
  • k (menyusut): {256, 512, ..., 4096}

Output dan Penyimpanan

Penyapuan penuh menghasilkan tabel latensi .pbtxt, yang siap digunakan oleh interpolator.

Interpolator

Interpolator adalah komponen compiler yang menggunakan tabel yang dihasilkan untuk memperkirakan performa GEMM.

Alasan: Saturasi FLOPS

Tabel latensi yang dikumpulkan memungkinkan interpolator merekonstruksi FLOPS untuk setiap entri:

\[\text{FLOPS} = \frac{2 \times b \times m \times n \times k}{\text{runtime} }\]

Insight utamanya adalah FLOPS jenuh pada titik tertentu; artinya, hardware mencapai FLOPS puncak di luar bentuk matriks tertentu. Saturasi ini memungkinkan penggunaan metode ekstrapolasi yang sama yang digunakan untuk kolektif.

Pencarian dan Interpolasi

Interpolator membangun ruang Euklidean 4D dari data tabel. Untuk memberikan perkiraan performa, model ini melakukan interpolasi rata-rata berbobot dalam ruang 4D ini. Jika tidak ada tabel untuk jenis data tertentu, setiap dimensi dinormalisasi ke jumlah byte sebagai heuristik.

Model Biaya Analisis - DCN

Model Biaya Kolektif Kurva-S

Model Kurva S adalah model batas atas jaringan yang sepenuhnya analitis.

Ringkasan

Model ini dirancang untuk memperkirakan performa operasi kolektif berdasarkan sekumpulan properti jaringan tetap.

Input Model

Model ini memerlukan dua kategori input:

  1. Properti Jaringan Tetap (Ditetapkan Pengguna):

    • Overhead peluncuran kolektif
    • Kecepatan NIC
    • RTT (waktu round trip)

    Secara default, XLA otomatis mendeteksi platform dan menggunakan nilai untuk arsitektur yang paling umum. Properti ini dapat dikonfigurasi oleh pengguna. Lihat bagian Penyesuaian untuk mengetahui detailnya.

  2. Input Per-Kolektif:

    • Jenis kolektif (misalnya, AllGather, ReduceScatter)
    • Ukuran transfer
    • Jumlah node yang terlibat dalam komunikasi

Integrasi

Model kurva S diintegrasikan ke dalam XLA:GPU dan digunakan di Hopper, dan Blackwell.

Model Biaya Analisis - Penggabungan

Untuk kernel lainnya, kami mengandalkan model biaya performa GPU untuk memperkirakan runtime yang tepat. Anda dapat membaca selengkapnya di sini.

Tuning mesin

Model kurva S dapat disesuaikan dengan mengeluarkan flag XLA yang tepat. Konfigurasi default sudah cukup baik dalam sebagian besar kasus, tetapi kontrol model ditampilkan dalam kasus lainnya.

export NIC_SPEED_GBPS=... # NIC speed per GPU in Gigabytes
export GPUS_PER_NODE=... # Num of GPUs per cluster interconnected with fast network (e.g. NVLINK)
export XLA_FLAGS=--xla_gpu_analytical_latency_estimator_options="nic_speed_gbps=$NIC_SPEED_GBPS,gpus_per_node=$GPUS_PER_NODE"