Tesla хочет использовать кремний для машинного обучения в Додзё

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

По сути, система представляет собой массивный компонуемый суперкомпьютер, хотя в отличие от того, что мы видим на Топ-500, он построен на основе полностью индивидуальной архитектуры, которая включает в себя вычислительные, сетевые устройства и микросхемы ввода-вывода (I/O), а также архитектуру набора команд (ISA), подачу питания, компоновку и охлаждение. Все это было сделано с явной целью запуска специализированных, конкретных алгоритмов обучения машинному обучению в масштабе.

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

Однако он утверждал, что традиционные методы масштабирования распределенных рабочих нагрузок не смогли ускориться со скоростью, необходимой для удовлетворения требований машинного обучения. По сути, закон Мура не ограничивает его, равно как и системы, доступные для масштабного обучения ИИ/МО, а именно некоторая комбинация ЦП/ГП или, в более редких случаях, с использованием специальных ускорителей ИИ. 

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

Бутерброд с центром обработки данных

Поэтому, чтобы обойти ограничения, Венкатараманан и его команда начали с нуля.

«Прямо во время моего интервью с Илоном он спросил меня, что вы можете сделать, что отличается от процессоров и графических процессоров для ИИ. Я чувствую, что вся команда все еще отвечает на этот вопрос».

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

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

«Эта учебная плитка представляет собой беспрецедентную степень интеграции от компьютера к памяти, к подаче энергии, к связи, не требуя никаких дополнительных переключателей», — сказал Венкатараманан.

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

«Если сравнивать транзисторы по квадратным миллиметрам, это, вероятно, передний край всего, что существует», — сказал Венкатараманан.

Затем 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 позволяет объединять память в большие группы графических процессоров. Однако Венкатараманан утверждает, что это не только увеличивает сложность, но и приводит к задержке и снижению пропускной способности.

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

Software

Такая специализированная вычислительная архитектура требует специализированного стека программного обеспечения. Однако Венкатараманан и его команда понимали, что программируемость либо спасет, либо разрушит Додзё.

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

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

«Мы использовали PiTorch. Мы создали промежуточный слой, который помогает нам распараллелить аппаратное обеспечение под ним. Под всем этим находится скомпилированный код», — сказал он. «Это единственный способ создать программные стеки, которые можно адаптировать ко всем этим будущим рабочим нагрузкам».

Несмотря на упор на гибкость программного обеспечения, Венкатараманан отмечает, что платформа, которая в настоящее время работает в их лабораториях, на данный момент ограничена использованием Tesla.

«В первую очередь мы ориентируемся на наших внутренних клиентов», — сказал он. «Илон объявил, что со временем мы сделаем это доступным для исследователей, но у нас нет для этого временных рамок. ®