Tesla ingin membawa silikon pembelajaran mesin ke Dojo

Untuk memuaskan dahaga akan model AI dan pembelajaran mesin yang semakin besar, Tesla telah mengungkapkan banyak detail di Hot Chips 34 tentang arsitektur superkomputer kustom penuh mereka yang disebut Dojo.

Sistem ini pada dasarnya adalah superkomputer besar yang dapat dikomposisi, meskipun tidak seperti yang kita lihat di Top 500, itu dibangun dari arsitektur yang sepenuhnya khusus yang mencakup silikon komputasi, jaringan, dan input/output (I/O) hingga arsitektur set instruksi (ISA), pengiriman daya, pengemasan, dan pendinginan. Semua itu dilakukan dengan tujuan yang jelas untuk menjalankan algoritme pelatihan pembelajaran mesin khusus yang disesuaikan dalam skala besar.

“Pemrosesan data dunia nyata hanya dapat dilakukan melalui teknik pembelajaran mesin, baik itu pemrosesan bahasa alami, mengemudi di jalan yang dibuat untuk penglihatan manusia hingga robotika yang berinteraksi dengan lingkungan sehari-hari,” Ganesh Venkataramanan, direktur senior teknik perangkat keras di Tesla, mengatakan selama pidato utamanya.

Namun, dia berargumen bahwa metode tradisional untuk menskalakan beban kerja terdistribusi gagal berakselerasi pada tingkat yang diperlukan untuk memenuhi tuntutan pembelajaran mesin. Akibatnya, Hukum Moore tidak memotongnya dan sistem juga tidak tersedia untuk pelatihan AI/ML dalam skala besar, yaitu beberapa kombinasi CPU/GPU atau dalam keadaan yang lebih jarang dengan menggunakan akselerator AI khusus. 

“Biasanya kami membuat chip, kami meletakkannya di paket, paket menggunakan PCB, yang masuk ke sistem. Sistem masuk ke rak, ”kata Venkataramanan. Masalahnya adalah setiap kali data berpindah dari chip ke paket dan keluar dari paket, itu menimbulkan penalti latency dan bandwidth.

Sandwich pusat data

Maka untuk menyiasati keterbatasan tersebut, Venkataramanan dan timnya memulai kembali dari awal.

“Tepat dari wawancara saya dengan Elon, dia bertanya kepada saya apa yang dapat Anda lakukan yang berbeda dari CPU dan GPU untuk AI. Saya merasa seluruh tim masih menjawab pertanyaan itu.”

Hal ini mengarah pada pengembangan ubin pelatihan Dojo, kluster komputasi mandiri yang menempati kaki setengah kubik yang mampu menghasilkan 556 TFLOPS kinerja FP32 dalam paket berpendingin cairan 15kW.

Setiap ubin dilengkapi dengan 11GB SRAM dan terhubung melalui fabric 9TB/dtk menggunakan protokol transport khusus di seluruh tumpukan.

“Ubin pelatihan ini mewakili jumlah integrasi yang tak tertandingi dari komputer ke memori hingga pengiriman daya, ke komunikasi, tanpa memerlukan sakelar tambahan apa pun,” kata Venkataramanan.

Di jantung ubin pelatihan adalah Tesla D1, transistor mati 50 miliar, berdasarkan proses 7nm TSMC. Tesla mengatakan setiap D1 mampu menghasilkan kinerja FP22 32 TFLOPS pada TDP 400W. Namun, Tesla mencatat bahwa chip tersebut mampu menjalankan berbagai perhitungan floating point termasuk beberapa yang khusus.

“Jika Anda membandingkan transistor untuk milimeter persegi, ini mungkin merupakan titik awal dari apa pun yang ada di luar sana,” kata Venkataramanan.

Tesla kemudian mengambil 25 D1, membuangnya untuk die yang dikenal baik, dan kemudian mengemasnya menggunakan teknologi sistem-on-wafer TSMC untuk "mencapai integrasi komputasi dalam jumlah besar pada latensi sangat rendah dan bandwidth sangat tinggi," katanya.

Namun, desain sistem-on-wafer dan arsitektur yang ditumpuk secara vertikal menimbulkan tantangan dalam hal pengiriman daya.

Menurut Venkataramanan, kebanyakan akselerator saat ini menempatkan daya secara langsung bersebelahan dengan silikon. Dan meskipun terbukti, pendekatan ini berarti area akselerator yang luas harus didedikasikan untuk komponen tersebut, yang membuatnya tidak praktis untuk Dojo, jelasnya. Sebagai gantinya, Tesla merancang chip mereka untuk mengirimkan daya secara langsung melalui bagian bawah cetakan. 

Putuskan semuanya

“Kita dapat membangun seluruh pusat data atau seluruh bangunan dari petak pelatihan ini, tetapi petak pelatihan hanyalah bagian komputasi. Kami juga perlu memberinya makan,” kata Venkataramanan.

Untuk ini, Tesla juga mengembangkan Dojo Interface Processor (DIP), yang berfungsi sebagai jembatan antara CPU host dan prosesor pelatihan. DIP juga berfungsi sebagai sumber memori bandwidth tinggi (HBM) bersama dan sebagai NIC 400Gbit/dtk berkecepatan tinggi.

Setiap DIP memiliki 32 GB HBM dan hingga lima kartu ini dapat dihubungkan ke ubin pelatihan dengan kecepatan 900 GB/dtk dengan agregat 4.5 TB/dtk ke host dengan total 160 GB HBM per ubin.

Pasangan konfigurasi V1 Tesla dari ubin ini – atau 150 D1 mati – dalam susunan yang didukung empat CPU host yang masing-masing dilengkapi dengan lima kartu DIP untuk mencapai kinerja exaflop BF16 atau CFP8 yang diklaim.

Secara keseluruhan, Venkataramanan mengatakan arsitekturnya – detail secara mendalam di sini by Platform Selanjutnya – memungkinkan Tesla mengatasi keterbatasan yang terkait dengan akselerator tradisional seperti Nvidia dan AMD.

“Cara kerja akselerator tradisional, biasanya Anda mencoba memasukkan seluruh model ke dalam setiap akselerator. Replikasi, lalu alirkan data melalui masing-masing,” katanya. “Apa yang terjadi jika kita memiliki model yang lebih besar dan lebih besar? Akselerator ini bisa gagal karena kehabisan memori.”

Ini bukan masalah baru, katanya. NV-switch Nvidia misalnya memungkinkan memori untuk dikumpulkan di seluruh bank besar GPU. Namun, Venkataramanan berpendapat ini tidak hanya menambah kompleksitas, tetapi juga memperkenalkan latensi dan kompromi pada bandwidth.

“Kami memikirkan ini sejak awal. Ubin komputasi kami dan masing-masing cetakan dibuat untuk menyesuaikan model besar, ”kata Venkataramanan.

Perangkat lunak

Arsitektur komputasi khusus semacam itu menuntut tumpukan perangkat lunak khusus. Namun, Venkataramanan dan timnya menyadari bahwa programabilitas akan membuat atau menghancurkan Dojo.

“Kemudahan programabilitas untuk perangkat lunak mitra sangat penting ketika kami merancang sistem ini,” katanya. “Para peneliti tidak akan menunggu perangkat lunak Anda menulis kernel tulisan tangan untuk beradaptasi dengan algoritme baru yang ingin kami jalankan.”

Untuk melakukan ini, Tesla membuang ide menggunakan kernel, dan merancang arsitektur Dojo di sekitar compiler.

“Apa yang kami lakukan adalah kami menggunakan PiTorch. Kami membuat lapisan perantara, yang membantu kami memparalelkan untuk menskalakan perangkat keras di bawahnya. Di bawah semuanya adalah kode yang dikompilasi, ”katanya. “Ini adalah satu-satunya cara untuk membuat tumpukan perangkat lunak yang dapat disesuaikan dengan semua beban kerja di masa depan.”

Terlepas dari penekanan pada fleksibilitas perangkat lunak, Venkataramanan mencatat bahwa platform, yang saat ini berjalan di laboratorium mereka, terbatas pada penggunaan Tesla untuk saat ini.

“Kami fokus pada pelanggan internal kami terlebih dahulu,” katanya. “Elon telah mengumumkan bahwa seiring waktu, kami akan membuat ini tersedia untuk para peneliti, tetapi kami tidak memiliki kerangka waktu untuk itu. ®