i made a TPU from scratch — Transcript

Full transcript

1. 0:00Nvidia.
2. 0:00>> AI chips, Nvidia.
3. 0:02>> AI, Nvidia.
4. 0:03>> Nvidia chips.
5. 0:03>> Newcomers to the club like Nvidia, they
6. 0:05power that new AI revolution. Up 5.4%.
7. 0:08>> It's up 55%.
8. 0:09>> Look, Nvidia made people rich like no
9. 0:11stock I've ever seen.
10. 0:13>> And this
11. 0:15is the backbone behind the whole thing.
12. 0:17The GPU.
13. 0:18Okay, maybe I lied. What's in front of
14. 0:20me is actually an FPGA. Unlike standard
15. 0:23GPUs and other chips that are
16. 0:24manufactured to the same design, this
17. 0:27chip is meant to be reprogrammed. And on
18. 0:29this board, I've programmed a TPU.
19. 0:32You see, here's the goal. I want to
20. 0:34create a small AI model on my computer
21. 0:36that can do something super cheap. And
22. 0:37in this case, I've [music] created an
23. 0:39art predictor, where the model will try
24. 0:41to guess what you're drawing as you're
25. 0:42drawing it. Now, the reason I went with
26. 0:44a project of this scope is due to
27. 0:46hardware limitations, which brings us
28. 0:47back to Nvidia.
29. 0:48>> It's about it's accelerated computing.
30. 0:49People don't give that enough credit.
31. 0:51And
32. 0:52I'm going to use it. We're number one.
33. 0:54We're number one.
34. 0:55>> You see, the strength of an AI model
35. 0:57comes from the hardware that supports
36. 0:59it. In 2020, OpenAI released GPT-3,
37. 1:02which was the first model to really make
38. 1:03public headway. Nowadays, we can look
39. 1:05back and see how out of date these
40. 1:06models seem, but even back then, when we
41. 1:08looked at the hardware specs, it's crazy
42. 1:10how many resources were needed to
43. 1:11support it. GPT-3 needed 350 GB of
44. 1:15memory to hold the 175 billion
45. 1:17parameters the model needed. In order to
46. 1:19support this, OpenAI needed to share the
47. 1:21model across eight GPUs, which alone
48. 1:24cost it about $200,000 of hardware.
49. 1:27Which also explains why memory stock
50. 1:29skyrocketed this year.
51. 1:30>> This is the biggest gainer in the S&P
52. 1:31500, up 19.3%. The market cap now over a
53. 1:35trillion dollars.
54. 1:36>> If we compare OpenAI's model to the size
55. 1:38of mine, it doesn't even show on the
56. 1:40charts. This chip in front of me has 512
57. 1:43KB of RAM and uses only 2 W of energy.
58. 1:46Yet, my thing still works and at 1/1000
59. 1:49of the cost.
60. 1:50Here's how. On the PC side, we'll have a
61. 1:5216 by 16 LED grid in which we can start
62. 1:55drawing. The plan is that as we draw,
63. 1:57the PC will take a snapshot of our grid
64. 1:59as a matrix. It will embed those
65. 2:01matrices into 64-dimensional token
66. 2:03vectors and add positional encoding. So,
67. 2:05each of the 16 rows becomes a token the
68. 2:08model can reason about. That 16 by 64
69. 2:10matrix then gets quantized to Q8.8
70. 2:13fixed-point format and gets sent out to
71. 2:15our chip, which is packed into a UART
72. 2:17frame with a checksum. The FPGA runs the
73. 2:19attention block on it, projecting the
74. 2:21tokens into queries, keys, and values,
75. 2:23computing the attention scores, applying
76. 2:25soft mats and hardware, and producing
77. 2:27the weighted output. It then sends the
78. 2:28resulting 16 by 64 matrix back over
79. 2:31UART. The PC picks up from there. It
80. 2:33adds a residual connection, applies a
81. 2:35layer norm, and the feed-forward
82. 2:36network. After that, it runs a final
83. 2:38output projection and a sigmoid, giving
84. 2:40a probability for each of the 256
85. 2:42pixels. Threshold those, and you have
86. 2:44the completed pattern, which gets drawn
87. 2:46back onto the grid. So, the loop is:
88. 2:48draw, snapshot, embed, FPGA does the
89. 2:51attention, PC finishes the layer,
90. 2:53predict, and display.
91. 2:55But, there is a bottleneck here, and
92. 2:58it's this:
93. 2:59over and over and over, the same
94. 3:02operation a few million times, all in
95. 3:05sequence.
96. 3:06So, if you step back, over 95% of this
97. 3:09is matrix multiplication, which is the
98. 3:11same expensive thing just repeated. A
99. 3:14normal CPU is going to chew through all
100. 3:16of this one by one, since it's a
101. 3:17general-purpose chip being asked to do
102. 3:19one thing a few million times. So, we
103. 3:21got to fix this on the hardware level,
104. 3:24and that's where we bring out the TPU.
105. 3:27Our chip provides us with configurable
106. 3:29logic, dedicated math blocks, and RAM.
107. 3:31On the chip itself, we'll need to
108. 3:33perform these operations.
109. 3:35First, the three projections. Because
110. 3:37the training gives us concrete embedding
111. 3:39matrices that won't change, we'll
112. 3:41actually have those pre-filled in our
113. 3:42chip's RAM. But, as we draw on the grid,
114. 3:45a snapshot will have to be transmitted
115. 3:46live. So, we'll need to make the
116. 3:48distinction as to what's live and what's
117. 3:50static. Let's take a look at the size of
118. 3:52our biggest operation. What we've got is
119. 3:54a 16 by 64 matrix times a 64 by 64
120. 3:57matrix. Perfect. Let's just design a
121. 4:00multiplication unit that can, I'm at, do
122. 4:03that.
123. 4:05We can't. See, there's an important
124. 4:07distinction to note on FPGAs and chips
125. 4:09in general, and that is the difference
126. 4:10between CLBs and DSPs.
127. 4:13An FPGA is mostly made up of
128. 4:15configurable logic blocks, which gives a
129. 4:17reprogrammable aspect to this chip.
130. 4:19These are made with lookup tables and
131. 4:21flip flops that you configure into
132. 4:22whatever circuit you need when you write
133. 4:24Verilog code. These can mimic anything,
134. 4:26and we've got about 20,000 of them on
135. 4:27the board. A DSP, however, is a
136. 4:30specialized non-programmable,
137. 4:31purpose-built silicon that does one
138. 4:33thing: multiply two numbers and add the
139. 4:35result to a running total, one cycle
140. 4:38every cycle, and fast. You can build a
141. 4:40multiplier out of CLBs, but a 16-bit one
142. 4:43eats around 200 lookup tables, and it's
143. 4:45slow because you're chaining logic
144. 4:47together, and the signal has to ripple
145. 4:49through all of the board, which might
146. 4:50not meet timing on our chip. And since
147. 4:52multiply-accumulate is the only
148. 4:54operation we'll really need, the DSPs
149. 4:56are the obvious choice. It's a The Basis
150. 4:593 only has 90 of them. For a dot
151. 5:01product, we need one DSP per output
152. 5:03element because each one is its own
153. 5:05running total. And if we go spec out how
154. 5:07many our biggest operation needed, that
155. 5:09would be 1,024 DSPs. So, how big can we
156. 5:12really make our matrix multiplier? 8 by
157. 5:158, as that will just use [music] 64
158. 5:17DSPs. Eight rows, eight columns, one DSP
159. 5:19per cell, all 64 running at once, which
160. 5:22gives us 26 to spare.
161. 5:24Not that it matters.
162. 5:25This works out nicely because every
163. 5:27matrix we're dealing with is a multiple
164. 5:29of eight. So, we can just chop them into
165. 5:308 by 8 tiles and feed them through one
166. 5:32at a time with nothing left over.
167. 5:3416 by 64 * 64 by 64 becomes two tiles
168. 5:38down, eight across, eight deep for 128
169. 5:40tiles run back-to-back through the same
170. 5:42array.
171. 5:43Now, here's the point when we make the
172. 5:44distinction between a GPU and a TPU. Try
173. 5:47to find the bottleneck in this
174. 5:48operation.
175. 5:51See, that's one cycle. Going out to
176. 5:53memory to fetch the values takes time,
177. 5:56and you're doing that fetch for every
178. 5:57single multiply. So, you end up with a
179. 5:59multiplier sitting idle most of the time
180. 6:01waiting on data that's slower to arrive
181. 6:03than it is to use. In standard GPU
182. 6:05architecture, it's done this way. We
183. 6:07multiply and fetch, but in TPU
184. 6:09architecture, we implement systolic
185. 6:11methodology in hopes of eliminating this
186. 6:13bottleneck.
187. 6:14The idea is that instead of every unit
188. 6:17fetching its own data, we fetch once and
189. 6:19let the data flow in this offset
190. 6:21pattern. We wire the multiply accumulate
191. 6:23units into a grid where each one is
192. 6:25connected to its neighbors. A value gets
193. 6:27read from memory once, enters the edge
194. 6:28of the array, and then gets passed hand
195. 6:30to hand from one unit to the next every
196. 6:32cycle one step further in. So, every
197. 6:35unit is doing a multiply every single
198. 6:37cycle, and the only memory traffic is at
199. 6:39the edges. The data watts through the
200. 6:41array, and by the time it comes out the
201. 6:42other side, all the multiply accumulates
202. 6:44are done.
203. 6:46So, when we apply it to our 8 by 8
204. 6:47multiply accumulator, it goes like this.
205. 6:50We load eight values of A and eight of B
206. 6:52at the edges of the array, and then all
207. 6:5364 units run for 22 cycles off that
208. 6:56single load. We're not fetching per
209. 6:58multiply, we're fetching per tile. So,
210. 7:00each chunk still gets that flow-through
211. 7:02benefit. The tiling doesn't break the
212. 7:04systolic idea. It just means we do it
213. 7:06128 times in a row instead of once on a
214. 7:08giant array. Where it does cost us is
215. 7:11the seams between tiles. As every time
216. 7:13we finish one and start the next, we go
217. 7:15back to memory for the next chunk. And
218. 7:16that load takes 64 cycles, which is
219. 7:19longer than the 22 cycles of actual
220. 7:21compute time. But, the good thing is
221. 7:23while the systolic array is computing
222. 7:25and the data's flowing through it, we
223. 7:26can actually start loading the next tile
224. 7:28in at the same time. Two buffers instead
225. 7:30of one. The array works out of one while
226. 7:32the controller fills the other. Then
227. 7:34they swap. The fetch hides underneath
228. 7:36the compute instead of stacking on top
229. 7:38of it. So, we lose no time between
230. 7:40tiles. This principle is called double
231. 7:42buffering, and it's what a real TPU
232. 7:43does. So, how much time did we really
233. 7:45save here? Well, let's look at the math.
234. 7:47Each 8 by 8 tile takes 64 cycles to
235. 7:50load, and with the double-buffering,
236. 7:52that's the entire cost. The 22 cycles of
237. 7:54compute hiding underneath it. One layer
238. 7:57is 576 tiles * 64 cycles is about 37,000
239. 8:02cycles. At 100 MHz, which is what the
240. 8:04board operates at, that's 0.37 ms. Six
241. 8:09matrix multiplications, half a million
242. 8:11multiply accumulates in about a third of
243. 8:13a second.
244. 8:14That's basically all the matrix math
245. 8:16done, and there's two parts left to do.
246. 8:18Number one is the softmax, which gives
247. 8:20probabilities to our attention tokens,
248. 8:22and two is the train the model.
249. 8:25The word softmax might sound complex,
250. 8:27but all it does is turn a row of raw
251. 8:29scores into probabilities. This is why
252. 8:31every input can produce a unique output
253. 8:33on something like ChatGPT.
254. 8:35Our attention scores come out as
255. 8:37arbitrary numbers, and what we need is
256. 8:39how much should this token pay attention
257. 8:41to that one as a fraction. Softmax takes
258. 8:44that row and squashes it, so
259. 8:45everything's positive and the whole row
260. 8:47sums to one. The way you do it is you
261. 8:49would exponentiate every score, then
262. 8:51divide each one by the total. But,
263. 8:53exponentials and divisions are two of
264. 8:55the most expensive things you can ask
265. 8:56the hardware to do. A divider alone is a
266. 8:58huge block of logic, and it's slow.
267. 9:01Where a multiply is one cycle, a
268. 9:02division takes many. Doing that for
269. 9:05every element in the matrix on a chip
270. 9:07that only has a 90 multipliers to begin
271. 9:09with is not happening.
272. 9:11So, we don't. There's a paper from 2021
273. 9:14that gets around it entirely using
274. 9:15lookup tables. Not to be confused with
275. 9:18the ones we saw earlier that are baked
276. 9:19into the FPGA.
277. 9:21The idea is that instead of calculating
278. 9:23the exponential, you precompute the
279. 9:25answers once and store them in a table.
280. 9:28Then at runtime, you just look it up.
281. 9:30Same for division. Rather than dividing,
282. 9:32you use the exponential and the sum as
283. 9:35coordinates into a second
284. 9:36two-dimensional table that hands you the
285. 9:39answer directly.
286. 9:40So, the whole thing becomes two table
287. 9:42reads and an adder, and both tables are
288. 9:44small enough to fit in a chip's general
289. 9:46purpose logic. So, it costs us zero RAM.
290. 9:49Since we're approximating, we'll be off
291. 9:51by about 0.9% on average. But, for
292. 9:54attention, what matters is which token
293. 9:56wins and not the exact decimal. And it
294. 9:58gets that right about 96% of the time.
295. 10:00After implementing the soft max, along
296. 10:02with a few other tweaks like writing the
297. 10:04UART driver, we're done with the
298. 10:06hardware.
299. 10:07Now, we actually have to train our
300. 10:09model. This spans two steps. Number one,
301. 10:12gather the data set. In our case, that's
302. 10:15like really simple, because we're
303. 10:16basically just doing pixel art. So, I
304. 10:18can literally just draw these out by
305. 10:19hand.
306. 10:20We've got 69 different pattern types:
307. 10:22hearts, skulls, Space Invaders, chess
308. 10:24pieces, letters, numbers, arrows, mazes,
309. 10:27a yin yang. Then for each one, it makes
310. 10:29six variants at slightly different
311. 10:30positions, plus noisy copies. So, the
312. 10:32model just can't memorize one heart.
313. 10:35That gives us 1,200 training examples. I
314. 10:37totally made 10,000, since it cost me
315. 10:39nothing. More data is free here, but
316. 10:41that's actually where we hit the first
317. 10:43wall in this project. See, every pattern
318. 10:45type I add [music] makes the job harder.
319. 10:47The model has to learn 69 different
320. 10:49shapes with the same set of weights, and
321. 10:50how much it can actually learn is capped
322. 10:52by how big the model is, and how big the
323. 10:54model can be is ultimately capped by the
324. 10:56board, which we know is not the best.
325. 10:59Every time you double the embedding
326. 11:00dimension, the attention weights
327. 11:01quadruple, because they're square
328. 11:03matrices, D by D.
329. 11:05At an embedding dimension of 64, which
330. 11:07is what we use, the whole model is about
331. 11:0969K parameters, which is 138 kilobytes
332. 11:12in our fixed point format. That fits on
333. 11:14our board, which has 225 kilobytes of
334. 11:16block RAM. Bump it to 128, now you're
335. 11:19roughly looking at 262K parameters,
336. 11:21which is over half a megabyte. The board
337. 11:24physically cannot hold that. Not that
338. 11:26it'd be slow or tight, but the memory
339. 11:28physically does not exist on the chip.
340. 11:30So, for this project, we're staying at
341. 11:32an embedding dimension of 64, which is
342. 11:34the biggest I could go with and still
343. 11:35have it fit. Loss is 0.11 at 95.8%
344. 11:39accuracy.
345. 11:41But, there's a trade-off buried in here
346. 11:42that I think is the actually interesting
347. 11:44part. I could have trained on way fewer
348. 11:46shapes, like just hearts and smileys,
349. 11:48and gotten near perfect accuracy with a
350. 11:50tiny model. But, then it's not really
351. 11:52learning anything. It's just memorizing
352. 11:53two pictures. Draw anything else and
353. 11:55it's not nothing. Throwing 69 different
354. 11:58shapes at it forces it to actually learn
355. 12:00structure, symmetry, edges, how shapes
356. 12:02close up, instead of memorizing specific
357. 12:04patterns, which is what I want. Because
358. 12:06when I draw something freehand, it's not
359. 12:08going to look nice and look like one of
360. 12:09my training examples.
361. 12:11But, that costs capacity. The more
362. 12:13variety you ask for, the bigger the
363. 12:14model needs to be to hold it. So, you're
364. 12:17getting squeezed from both ends. Too few
365. 12:19shapes and it memorizes, too many and
366. 12:21you need a model the board can't fit.
367. 12:23Step two, train the model.
368. 12:26Luckily for us, though, training is the
369. 12:28easy part, since PyTorch does it all. We
370. 12:31describe the shape of the network, hand
371. 12:32it the data, and it works out every
372. 12:34adjustment for us.
373. 12:35>> [music]
374. 12:35>> One decision worth mentioning, at 95.8%
375. 12:38accuracy, I wasn't happy with a single
376. 12:40attention layer, so I ran two, which
377. 12:42means two full sets of weight matrices
378. 12:44instead of one. After training, all we
379. 12:46walk away with is just matrices,
380. 12:49embedding, QKV, output projection, feed
381. 12:52forward, and one last one to turn it
382. 12:54back into pixels.
383. 12:58And with that, the project's done. I
384. 13:01finished my code, got my board up and
385. 13:02running, and connected my PC with the
386. 13:04weights loaded to see how well this
387. 13:06actually works. For a quick test, if I
388. 13:08click any tile in our predictor, the
389. 13:10UART LED should flash, which means it's
390. 13:12set up to test. So, now, moment of
391. 13:14truth. Let's see how good this TPU
392. 13:17actually is.
393. 13:21If I draw like this, we can see that it
394. 13:22gets it fairly correct. 96% accuracy
395. 13:25seems to get us not the healthiest
396. 13:27heart, but to be honest, I'm quite
397. 13:29surprised with how well this turned out
398. 13:30with all the other drawings. This tiny
399. 13:32board is reading my input, running the
400. 13:34two-layer inference, and is guessing
401. 13:36what I'm drawing quite well. And yeah,
402. 13:38with some shapes, it doesn't provide
403. 13:40exactly what I'm looking for, but it's
404. 13:42still giving a really good approximation
405. 13:44based on the fact that I'm drawing like
406. 13:4630% of the actual shape. And honestly,
407. 13:49it does this better than me in some of
408. 13:50these.
409. 13:51I would not have known this was a boat,
410. 13:53but that's probably just because I'm a
411. 13:55bad drawer. For comparison's sake, I
412. 13:57decided to run the same script, but just
413. 13:59using the specs of my PC. I used PyTorch
414. 14:02again, not just to train the data, but
415. 14:03to help implement the attention pipeline
416. 14:05solely on my CPU.
417. 14:07And I think the best part of this whole
418. 14:09thing is that my CPU is a $400 chip. And
419. 14:13while it beats my FPGA in speed, it
420. 14:15would have been beat if we took the wire
421. 14:17out of the equation and didn't have to
422. 14:18spend 21 milliseconds on the UART
423. 14:21transfer.
424. 14:23Intel, AMD, and other companies spend
425. 14:25decades and billions of dollars to get
426. 14:27to this point. And this is a $60 FPGA I
427. 14:30programmed at my desk over a few months.
428. 14:33And I think in some aspects, it wins.
429. 14:35The thing behind the trillion-dollar
430. 14:37companies, the thing on every news
431. 14:39channel right now, the thing that sounds
432. 14:41like it should be out of reach, turns
433. 14:43out you can just make one.
434. 14:56>> Mhm.