تسلا می خواهد سیلیکون یادگیری ماشین را به دوجو ببرد

تسلا برای رفع تشنگی مدل‌های هوش مصنوعی و یادگیری ماشینی بزرگ‌تر، جزئیات زیادی را در Hot Chips 34 در مورد معماری ابرکامپیوتری کاملا سفارشی خود به نام Dojo فاش کرده است.

این سیستم اساساً یک ابررایانه بزرگ قابل ترکیب است، اگرچه برخلاف آنچه در آن می‌بینیم بالا 500, it’s built from an entirely custom architecture that spans the compute, networking, and input/output (I/O) silicon to instruction set architecture (ISA), power delivery, packaging, and cooling. All of it was done with the express purpose of running tailored, specific machine learning training algorithms at scale.

“Real world data processing is only feasible through machine learning techniques, be it natural-language processing, driving in streets that are made for human vision to robotics interfacing with the everyday environment,” Ganesh Venkataramanan, senior director of hardware engineering at Tesla, said during his keynote speech.

However, he argued that traditional methods for scaling distributed workloads have failed to accelerate at the rate necessary to keep up with machine learning’s demands. In effect, Moore’s Law is not cutting it and neither are the systems available for AI/ML training at scale, namely some combination of CPU/GPU or in rarer circumstances by using speciality AI accelerators. 

“Traditionally we build chips, we put them on packages, packages go on PCBs, which go into systems. Systems go into racks,” said Venkataramanan. The problem is each time data moves from the chip to the package and off the package, it incurs a latency and bandwidth penalty.

ساندویچ دیتاسنتر

بنابراین برای دور زدن محدودیت ها، Venkataramanan و تیمش از ابتدا شروع کردند.

“Right from my interview with Elon, he asked me what can you do that is different from CPUs and GPUs for AI. I feel that the whole team is still answering that question.”

این منجر به توسعه کاشی آموزشی Dojo شد، یک خوشه محاسباتی مستقل که نیم فوت مکعب را اشغال می‌کند و قادر به عملکرد 556 TFLOPS FP32 در یک بسته 15 کیلوواتی خنک‌کننده مایع است.

هر کاشی مجهز به 11 گیگابایت SRAM است و با استفاده از یک پروتکل حمل و نقل سفارشی در کل پشته از طریق یک پارچه 9 ترابایت بر ثانیه متصل می شود.

“This training tile represents unparalleled amounts of integration from computer to memory to power delivery, to communication, without requiring any additional switches,” Venkataramanan said.

At the heart of the training tile is Tesla’s D1, a 50 billion transistor die, based on TSMC’s 7nm process. Tesla says each D1 is capable of 22 TFLOPS of FP32 performance at a TDP of 400W. However, Tesla notes that the chip is capable of running a wide range of floating point calculations including a few custom ones.

“If you compare transistors for millimeter square, this is probably the bleeding edge of anything which is out there,” Venkataramanan said.

Tesla then took 25 D1s, binned them for known good dies, and then packaged them using TSMC’s system-on-wafer technology to “achieve a huge amount of compute integration at very low latency and very-high bandwidth,” he said.

با این حال، طراحی سیستم روی ویفر و معماری انباشته عمودی چالش هایی را در هنگام تحویل نیرو به وجود آورد.

به گفته Venkataramanan، امروزه بیشتر شتاب دهنده ها نیرو را مستقیماً در مجاورت سیلیکون قرار می دهند. و در حالی که این رویکرد ثابت شده است، به این معنی است که یک منطقه بزرگ از شتاب دهنده باید به آن اجزا اختصاص داده شود، که آن را برای Dojo غیر عملی می کند. درعوض، تسلا تراشه‌های خود را طوری طراحی کرد که نیرو را مستقیماً از پایین قالب منتقل کند. 

همه اش رو بذار کنار هم

“We could build an entire datacenter or an entire building out of this training tile, but the training tile is just the compute portion. We also need to feed it,” Venkataramanan said.

برای این کار، تسلا همچنین پردازنده رابط Dojo (DIP) را توسعه داد که به عنوان پلی بین CPU میزبان و پردازنده های آموزشی عمل می کند. DIP همچنین به عنوان منبع حافظه مشترک با پهنای باند بالا (HBM) و به عنوان یک NIC پرسرعت 400Gbit/sec عمل می کند.

هر DIP دارای 32 گیگابایت HBM است و حداکثر پنج عدد از این کارت ها را می توان به یک کاشی آموزشی با سرعت 900 گیگابایت بر ثانیه برای مجموع 4.5 ترابایت بر ثانیه به میزبان و در مجموع 160 گیگابایت HBM در هر کاشی متصل کرد.

Tesla’s V1 configuration pairs of these tiles – or 150 D1 dies – in array supported four host CPUs each equipped with five DIP cards to achieve a claimed exaflop of BF16 or CFP8 performance.

در کنار هم، ونکاتارامان می گوید معماری - مفصل در عمق اینجا by پلت فرم بعدی - تسلا را قادر می سازد بر محدودیت های مربوط به شتاب دهنده های سنتی مانند Nvidia و AMD غلبه کند.

“How traditional accelerators work, typically you try to fit an entire model into each accelerator. Replicate it, and then flow the data through each of them,” he said. “What happens if we have bigger and bigger models? These accelerators can fall flat because they run out of memory.”

This isn’t a new problem, he noted. Nvidia’s NV-switch for example enables memory to be pooled across large banks of GPUs. However, Venkataramanan argues this not only adds complexity, but introduces latency and compromises on bandwidth.

“We thought about this right from the get go. Our compute tiles and each of the dies were made for fitting big models,” Venkataramanan said.

نرم افزار

چنین معماری محاسباتی تخصصی نیازمند یک پشته نرم افزار تخصصی است. با این حال، ونکاتارامانان و تیمش دریافتند که قابلیت برنامه‌ریزی، دوجو را یا می‌سازد یا خراب می‌کند.

“Ease of programmability for software counterparts is paramount when we design these systems,” he said. “Researchers won’t wait for your software folks to write a handwritten kernel for adapting to a new algorithm that we want to run.”

To do this, Tesla ditched the idea of using kernels, and designed Dojo’s architecture around compilers.

“What we did was we used PiTorch. We created an intermediate layer, which helps us parallelize to scale out hardware beneath it. Underneath everything is compiled code,” he said. “This is the only way to create software stacks that are adaptable to all those future workloads.”

با وجود تاکید بر انعطاف پذیری نرم افزار، Venkataramanan خاطرنشان می کند که این پلتفرم که در حال حاضر در آزمایشگاه های آنها در حال اجرا است، در حال حاضر به استفاده از تسلا محدود شده است.

“We are focused on our internal customers first,” he said. “Elon has made it public that over time, we will make this available to researchers, but we don’t have a time frame for that. ®