Tesla хоче взяти кремній для машинного навчання в Dojo

Щоб втамувати спрагу все більших моделей штучного інтелекту та машинного навчання, Tesla розкрила безліч деталей на Hot Chips 34 про свою повністю індивідуальну суперкомп’ютерну архітектуру під назвою Dojo.

Система, по суті, є масивним складним суперкомп’ютером, хоча й на відміну від того, що ми бачимо на Кращі 500, він побудований на основі повністю спеціальної архітектури, яка охоплює обчислення, мережу та введення/виведення (I/O) з архітектурою набору інструкцій (ISA), доставку живлення, упаковку та охолодження. Усе це було зроблено з явною метою запуску спеціальних навчальних алгоритмів машинного навчання в масштабі.

«Обробка даних у реальному світі можлива лише за допомогою методів машинного навчання, будь то обробка природної мови, їзда вулицями, створеними для людського зору, і роботизована техніка, яка взаємодіє з повсякденним середовищем», — сказав Ганеш Венкатараманан, старший директор апаратного забезпечення Tesla. під час його основної промови.

Однак він стверджував, що традиційні методи масштабування розподілених робочих навантажень не змогли прискоритися зі швидкістю, необхідною для виконання вимог машинного навчання. По суті, закон Мура не скорочує його, як і системи, доступні для навчання AI/ML у масштабі, а саме деяка комбінація CPU/GPU або, у рідкісних випадках, за допомогою спеціальних прискорювачів AI. 

«Традиційно ми створюємо чіпи, розміщуємо їх на упаковках, упаковки йдуть на друковані плати, які надходять у системи. Системи встановлюються в стійки», — сказав Венкатараманан. Проблема полягає в тому, що кожного разу, коли дані переміщуються з чіпа в пакет і з пакета, це спричиняє затримку та штраф пропускної здатності.

Сендвіч центру обробки даних

Тому, щоб обійти обмеження, Венкатараманан і його команда почали все з нуля.

«Одразу під час мого інтерв’ю з Ілоном він запитав мене, чим ти можеш відрізнятися від процесорів і графічних процесорів для ШІ. Я відчуваю, що вся команда все ще відповідає на це запитання».

Це призвело до розробки навчальної плитки Dojo, автономного обчислювального кластера, що займає половину кубічного фута, здатного забезпечити продуктивність 556 TFLOPS FP32 у корпусі з рідинним охолодженням потужністю 15 кВт.

Кожна плитка оснащена 11 ГБ пам’яті SRAM і з’єднана через структуру 9 ТБ/с за допомогою спеціального транспортного протоколу по всьому стеку.

«Ця навчальна плитка представляє неперевершений обсяг інтеграції від комп’ютера до пам’яті до живлення та зв’язку, не вимагаючи жодних додаткових комутаторів», — сказав Венкатараманан.

В основі навчальної плитки лежить Tesla D1, 50-мільярдний транзисторний кристал, заснований на 7-нм техпроцесі TSMC. Tesla каже, що кожен D1 здатний досягти 22 TFLOPS продуктивності FP32 при TDP 400 Вт. Однак Tesla зазначає, що чіп здатний запускати широкий спектр обчислень з плаваючою комою, включаючи кілька спеціальних.

«Якщо ви порівнюєте транзистори на квадратний міліметр, це, мабуть, кращий край усього, що існує», — сказав Венкатараманан.

Потім Tesla взяв 25 D1, об’єднав їх для завідомо справних матриць, а потім упакував їх за допомогою технології системи на пластині TSMC, щоб «досягти величезної кількості обчислювальної інтеграції з дуже низькою затримкою та дуже високою пропускною здатністю», – сказав він.

Однак дизайн «система на пластині» та архітектура з вертикальним штабелем створили проблеми, коли справа дійшла до доставки електроенергії.

За словами Венкатараманана, більшість сучасних прискорювачів розміщують потужність безпосередньо поруч із кремнієм. І хоча цей підхід перевірено, він означає, що велика площа прискорювача повинна бути виділена для цих компонентів, що робить його непрактичним для Dojo, пояснив він. Натомість Tesla розробила свої мікросхеми для передачі енергії безпосередньо через нижню частину матриці. 

Встановити все це разом

«Ми могли б побудувати цілий центр обробки даних або цілу будівлю з цієї навчальної плитки, але навчальна плитка — це лише обчислювальна частина. Нам також потрібно його годувати», — сказав Венкатараманан.

Для цього Tesla також розробила інтерфейсний процесор Dojo (DIP), який функціонує як міст між центральним процесором і навчальними процесорами. DIP також служить джерелом спільної пам’яті високої пропускної здатності (HBM) і високошвидкісною мережевою карткою 400 Гбіт/с.

Кожен DIP містить 32 ГБ пам’яті HBM, і до п’яти таких карт можна під’єднати до тренувальної плитки зі швидкістю 900 ГБ/с для загальної пропускної спроможності 4.5 ТБ/с для хосту, тобто 160 ГБ HBM на плитку.

Конфігураційні пари Tesla V1 із цих плиток – або 150 матриць D1 – у масиві підтримували чотири хост-ЦП, кожен із яких оснащено п’ятьма DIP-картами для досягнення заявленої продуктивності BF16 або CFP8 у екзафлопі.

За словами Венкатараманана, архітектура – ​​детальна тут глибоко by Наступна платформа – дозволяє Tesla подолати обмеження, пов’язані з традиційними прискорювачами, такими як Nvidia та AMD.

«Як працюють традиційні прискорювачі, зазвичай ви намагаєтеся вмістити цілу модель у кожен прискорювач. Відтворіть його, а потім передайте дані через кожен із них», — сказав він. «Що станеться, якщо ми матимемо все більші й більші моделі? Ці прискорювачі можуть вийти з ладу, тому що їм не вистачає пам’яті».

Це не нова проблема, зазначив він. NV-перемикач Nvidia, наприклад, дозволяє об’єднувати пам’ять у великі банки графічних процесорів. Однак Венкатараманан стверджує, що це не тільки ускладнює, але й створює затримку та знижує пропускну здатність.

«Ми думали про це з самого початку. Наші обчислювальні плитки та кожна матриця були створені для встановлення великих моделей», — сказав Венкатараманан.

Софтвер

Така спеціалізована обчислювальна архітектура вимагає спеціального програмного стека. Однак Венкатараманан і його команда визнали, що програмованість або допоможе, або зруйнує Dojo.

«Простота програмування аналогів програмного забезпечення має першорядне значення, коли ми розробляємо ці системи», — сказав він. «Дослідники не чекатимуть, поки ваші спеціалісти з програмного забезпечення напишуть рукописне ядро ​​для адаптації до нового алгоритму, який ми хочемо запустити».

Для цього Tesla відмовився від ідеї використання ядер і розробив архітектуру Dojo навколо компіляторів.

«Ми використовували PiTorch. Ми створили проміжний рівень, який допомагає нам розпаралелювати, щоб масштабувати обладнання під ним. Під усім скомпільований код», — сказав він. «Це єдиний спосіб створити стеки програмного забезпечення, які адаптуються до всіх цих майбутніх навантажень».

Незважаючи на наголос на гнучкості програмного забезпечення, Венкатараманан зазначає, що платформа, яка зараз працює в їхніх лабораторіях, на даний момент обмежена використанням Tesla.

«Насамперед ми зосереджені на наших внутрішніх клієнтах, — сказав він. «Ілон оприлюднив, що з часом ми зробимо це доступним для дослідників, але у нас немає часових рамок для цього. ®