Tesla muốn đưa silicon học máy vào Dojo

Để làm dịu cơn khát về các mô hình học máy và AI lớn hơn bao giờ hết, Tesla đã tiết lộ rất nhiều chi tiết tại Hot Chips 34 về kiến ​​trúc siêu máy tính hoàn toàn tùy chỉnh của họ có tên là Dojo.

Hệ thống này về cơ bản là một siêu máy tính có thể tổng hợp lớn, mặc dù không giống như những gì chúng ta thấy trên Top 500, nó được xây dựng từ một kiến ​​trúc hoàn toàn tùy chỉnh trải rộng từ silicon tính toán, kết nối mạng và đầu vào/đầu ra (I/O) đến kiến ​​trúc tập lệnh (ISA), phân phối điện, đóng gói và làm mát. Tất cả đều được thực hiện với mục đích rõ ràng là chạy các thuật toán đào tạo máy học cụ thể, phù hợp trên quy mô lớn.

Ganesh Venkataramanan, giám đốc cấp cao về kỹ thuật phần cứng của Tesla, cho biết: “Việc xử lý dữ liệu trong thế giới thực chỉ khả thi thông qua các kỹ thuật máy học, có thể là xử lý ngôn ngữ tự nhiên, lái xe trên những con phố được tạo ra cho tầm nhìn của con người cho đến việc robot giao tiếp với môi trường hàng ngày”. trong bài phát biểu quan trọng của mình.

Tuy nhiên, ông lập luận rằng các phương pháp truyền thống để mở rộng quy mô khối lượng công việc phân tán đã không thể tăng tốc ở tốc độ cần thiết để theo kịp nhu cầu của máy học. Trên thực tế, Định luật Moore không cắt giảm nó và các hệ thống cũng không có sẵn để đào tạo AI/ML trên quy mô lớn, cụ thể là sự kết hợp nào đó giữa CPU/GPU hoặc trong những trường hợp hiếm hơn bằng cách sử dụng các bộ tăng tốc AI đặc biệt. 

“Theo truyền thống, chúng tôi chế tạo chip, đặt chúng trên các gói, gói trên PCB rồi đi vào hệ thống. Hệ thống đi vào giá đỡ,” Venkataramanan nói. Vấn đề là mỗi khi dữ liệu di chuyển từ chip sang gói và ra khỏi gói, nó sẽ phải chịu một hình phạt về độ trễ và băng thông.

Một bánh sandwich trung tâm dữ liệu

Vì vậy, để vượt qua những hạn chế, Venkataramanan và nhóm của anh ấy đã bắt đầu lại từ đầu.

“Ngay từ cuộc phỏng vấn của tôi với Elon, anh ấy đã hỏi tôi rằng bạn có thể làm gì khác với CPU và GPU cho AI. Tôi cảm thấy cả đội vẫn đang trả lời câu hỏi đó.”

Điều này dẫn đến sự phát triển của gạch luyện tập Dojo, một cụm máy tính khép kín chiếm nửa foot khối có khả năng đạt 556 TFLOPS hiệu suất FP32 trong một gói làm mát bằng chất lỏng 15kW.

Mỗi ô được trang bị 11GB SRAM và được kết nối qua kết cấu 9TB / s bằng cách sử dụng giao thức truyền tải tùy chỉnh trong toàn bộ ngăn xếp.

Venkataramanan cho biết: “Ô đào tạo này thể hiện mức độ tích hợp tuyệt vời từ máy tính đến bộ nhớ, cung cấp năng lượng, giao tiếp mà không yêu cầu bất kỳ bộ chuyển mạch bổ sung nào”.

Tâm điểm của ô đào tạo là D1 của Tesla, khuôn bóng bán dẫn trị giá 50 tỷ, dựa trên quy trình 7nm của TSMC. Tesla cho biết mỗi D1 có khả năng đạt 22 TFLOPS hiệu suất FP32 ở mức TDP 400W. Tuy nhiên, Tesla lưu ý rằng con chip này có khả năng thực hiện nhiều phép tính dấu phẩy động, bao gồm một số phép tính tùy chỉnh.

Venkataramanan nói: “Nếu bạn so sánh các bóng bán dẫn theo từng milimet vuông, thì đây có lẽ là điểm vượt trội của bất kỳ thứ gì hiện có”.

Sau đó, Tesla đã lấy 25 chiếc D1, phân loại chúng để có những khuôn tốt đã biết và sau đó đóng gói chúng bằng công nghệ hệ thống trên tấm wafer của TSMC để “đạt được lượng tích hợp điện toán khổng lồ ở độ trễ rất thấp và băng thông rất cao,” ông nói.

Tuy nhiên, thiết kế system-on-wafer và kiến ​​trúc xếp chồng lên nhau theo chiều dọc đã đưa ra những thách thức khi phân phối điện.

Theo Venkataramanan, hầu hết các máy gia tốc ngày nay đều đặt nguồn điện tiếp giáp trực tiếp với silicon. Và trong khi đã được chứng minh, cách tiếp cận này có nghĩa là một khu vực lớn của máy gia tốc phải được dành riêng cho những thành phần đó, điều này khiến cho Dojo không thực tế, ông giải thích. Thay vào đó, Tesla đã thiết kế những con chip của họ để cung cấp điện năng trực tiếp từ đáy của khuôn. 

Để tất cả chúng cùng nhau

“Chúng tôi có thể xây dựng toàn bộ trung tâm dữ liệu hoặc toàn bộ tòa nhà từ ô đào tạo này, nhưng ô đào tạo chỉ là phần tính toán. Chúng tôi cũng cần cho nó ăn”, Venkataramanan nói.

Để làm được điều này, Tesla cũng đã phát triển Bộ xử lý Giao diện Dojo (DIP), có chức năng như một cầu nối giữa CPU chủ và các bộ xử lý đào tạo. DIP cũng đóng vai trò như một nguồn bộ nhớ băng thông cao được chia sẻ (HBM) và như một NIC tốc độ cao 400Gbit / giây.

Mỗi DIP có 32GB HBM và tối đa năm thẻ trong số này có thể được kết nối với ô đào tạo ở tốc độ 900GB / s với tổng cộng 4.5TB / s tới máy chủ để có tổng cộng 160GB HBM trên mỗi ô.

Các cặp cấu hình V1 của Tesla gồm các ô này – hoặc 150 khuôn D1 – xếp thành dãy hỗ trợ bốn CPU chủ, mỗi CPU được trang bị năm thẻ DIP để đạt được exaflop được công bố về hiệu suất BF16 hoặc CFP8.

Tổng hợp lại, Venkataramanan cho biết kiến ​​trúc - chi tiết sâu ở đây by Nền tảng tiếp theo - cho phép Tesla vượt qua những hạn chế liên quan đến các máy gia tốc truyền thống của Nvidia và AMD.

“Cách hoạt động của các máy gia tốc truyền thống, thông thường bạn cố gắng lắp toàn bộ mô hình vào mỗi máy gia tốc. Hãy sao chép nó và sau đó truyền dữ liệu qua từng cái,” ông nói. “Điều gì sẽ xảy ra nếu chúng ta có những mô hình ngày càng lớn hơn? Những máy gia tốc này có thể bị hỏng vì hết bộ nhớ.”

Đây không phải là một vấn đề mới, ông lưu ý. Ví dụ, NV-switch của Nvidia cho phép tập hợp bộ nhớ trên các dãy GPU lớn. Tuy nhiên, Venkataramanan lập luận rằng điều này không chỉ làm tăng thêm độ phức tạp mà còn gây ra độ trễ và ảnh hưởng đến băng thông.

“Chúng tôi đã nghĩ đến điều này ngay từ đầu. Venkataramanan cho biết: Các ô tính toán của chúng tôi và mỗi khuôn đều được chế tạo để lắp các mô hình lớn.

Phần mềm

Một kiến ​​trúc máy tính chuyên biệt như vậy đòi hỏi một ngăn xếp phần mềm chuyên dụng. Tuy nhiên, Venkataramanan và nhóm của ông nhận ra rằng khả năng lập trình sẽ tạo ra hoặc phá vỡ Dojo.

Ông nói: “Khả năng lập trình dễ dàng đối với các đối tác phần mềm là điều tối quan trọng khi chúng tôi thiết kế các hệ thống này. “Các nhà nghiên cứu sẽ không đợi người làm phần mềm của bạn viết hạt nhân viết tay để thích ứng với thuật toán mới mà chúng tôi muốn chạy.”

Để làm được điều này, Tesla đã từ bỏ ý tưởng sử dụng hạt nhân và thiết kế kiến ​​trúc của Dojo xung quanh các trình biên dịch.

“Những gì chúng tôi đã làm là sử dụng PiTorch. Chúng tôi đã tạo một lớp trung gian giúp chúng tôi song song hóa việc mở rộng quy mô phần cứng bên dưới nó. Bên dưới mọi thứ đều là mã được biên dịch,” ông nói. “Đây là cách duy nhất để tạo ra các kho phần mềm có khả năng thích ứng với tất cả khối lượng công việc trong tương lai.”

Mặc dù nhấn mạnh vào tính linh hoạt của phần mềm, Venkataramanan lưu ý rằng nền tảng, hiện đang chạy trong phòng thí nghiệm của họ, chỉ được Tesla sử dụng vào thời điểm hiện tại.

Ông nói: “Trước tiên, chúng tôi tập trung vào khách hàng nội bộ của mình. “Elon đã công khai rằng theo thời gian, chúng tôi sẽ cung cấp thông tin này cho các nhà nghiên cứu, nhưng chúng tôi không có khung thời gian cho việc đó. ®