Sự bùng nổ về cải thiện hiệu năng GPU
(hazyresearch.stanford.edu)- Trong bối cảnh chi phí tính toán AI ngày càng lớn, Hazy Research tổng kết rằng trọng tâm của tối ưu hiệu năng GPU là giữ cho tensor core của NVIDIA H100 luôn hoạt động không nghỉ
- H100 đạt 989 TFLOPs ở phép nhân ma trận half-precision, nhưng các phép toán thông thường chỉ khoảng 60 TFLOPs, nên ngay khi tensor core dừng lại, mức sử dụng giảm mạnh
- Để tiến gần hiệu năng tối đa, cần xử lý đồng thời WGMMA, bố trí shared memory, sinh địa chỉ và occupancy; nếu không có
wgmma.mma_async, microbenchmark chỉ đạt khoảng 63% đỉnh - DSL nhúng trong CUDA được công khai, ThunderKittens, bao bọc các phức tạp như swizzling và register layout bằng trừu tượng tile·vector, giúp đơn giản hóa việc viết kernel họ FlashAttention
- Kernel FlashAttention-2 forward cho H100 được viết trong khoảng 100 dòng, nhanh hơn FlashAttention-2 khoảng 30%, còn kernel Based linear attention chạy ở 215 TFLOPs
Các điều kiện quyết định hiệu năng H100
- AI sử dụng rất nhiều compute, và trong vài năm gần đây Hazy Research đã làm các công việc giúp AI dùng ít compute hơn hoặc chạy hiệu quả hơn trên lượng compute cho trước
- Ví dụ về tiết kiệm compute gồm Based, Monarch Mixer, H3, Hyena, S4
- Ví dụ về thực thi hiệu quả gồm FlashAttention, FlashAttention-2, FlashFFTConv
- Mục tiêu thực tiễn là tổng kết những gì học được khi làm GPU chạy nhanh hơn, và công khai ThunderKittens, một DSL nhúng trong CUDA giúp viết kernel nhanh
- Rộng hơn, bài viết bàn về việc hiểu phần cứng đã thay đổi cách nhìn về compute cho AI như thế nào
Cấu trúc và nút thắt của NVIDIA H100
- GPU H100 SXM được thảo luận dựa trên cấu hình sau
- 80GB HBM3, băng thông 3TB/s
- 50MB L2 cache, băng thông 12TB/s, được chia thành 2 phần 25MB trên toàn GPU và kết nối bằng crossbar
-
132 SM
- Mỗi SM có 256KB L1 cache, bao gồm tối đa 227KB shared memory, và tổng cộng có băng thông khoảng 33TB/s
- Phần cứng mới của Hopper, Tensor Memory Accelerator(TMA), đảm nhiệm sinh địa chỉ bất đồng bộ và fetch bộ nhớ
- Mỗi SM gồm 4 quadrant, mỗi quadrant có warp scheduler, 512 vector register, tensor core cho phép nhân ma trận và các lệnh dựng sẵn song song
- Mọi compute đều diễn ra trong SM, và phần lớn được xử lý trong register
- Điểm then chốt để đạt hiệu năng trên H100 là giữ tensor core luôn được fed dữ liệu
- H100 cung cấp 989 TFLOPs cho phép nhân ma trận half-precision, và khoảng 60 TFLOPs cho “các phép toán khác”
- Trong các cycle mà tensor core được dùng, mức sử dụng phần cứng đạt tối thiểu 94%
- Trong các cycle mà tensor core không được dùng, mức sử dụng tối đa chỉ dừng ở 6%
WGMMA: lệnh cần thiết nhưng khó dùng
- H100 có lệnh warp group matrix multiply accumulate là
wgmma.mma_async- Trong PTX là
wgmma.mma_async - Trong SASS là
HGMMA/IGMMA/QGMMA/BGMMA
- Trong PTX là
wmma.mma.sync,mma.synctrên các GPU trước đây là kiểu đồng bộ, trong đó một warp gồm 32 thread đưa dữ liệu vào tensor core rồi chờ kết quảwgmma.mma_asynccho phép 128 thread liên tiếp đồng bộ cộng tác trên tất cả quadrant của SM, và khởi chạy phép nhân ma trận bất đồng bộ trực tiếp từ shared memory- Các warp có thể làm việc khác bằng register trong khi phép nhân ma trận đang diễn ra
- Có thể chờ kết quả vào thời điểm mong muốn
- Trong microbenchmark, các lệnh này là cần thiết để khai thác toàn bộ compute của H100
- Nếu không dùng, GPU được quan sát là chỉ dừng ở khoảng 63% mức sử dụng đỉnh
- Có thể vì tensor core đòi hỏi pipeline phần cứng sâu ngay cả với tài nguyên cục bộ
- Khó khăn lớn nhất là sự phức tạp của memory layout
- Unswizzled shared memory layout có coalescing rất kém, đòi hỏi nhiều băng thông L2
- Swizzled layout có tài liệu sai, nên mất thời gian để hiểu
- Swizzled layout có vẻ chỉ hoạt động với một số shape ma trận nhất định, và không tương thích tốt với các tính năng khác của
wgmma.mma_async - Phần cứng có thể thực hiện sub-matrix transpose trên đường tới tensor core, nhưng chỉ khi layout không phải swizzled
- Trong các kernel như FlashAttention, TMA và L2 cache đủ nhanh để che giấu vấn đề này ở một mức độ nào đó
- Để tận dụng hoàn toàn phần cứng, cần coalescing memory request và tránh bank conflict, nên việc kiểm soát layout rất quan trọng
Shared memory và bank conflict
- Single-access latency của shared memory có vẻ khoảng 30 cycles, trong thời gian đó tensor core của SM gần như có thể thực hiện hai phép nhân ma trận vuông 32x32
- Trong các công trình trước đây như FlashAttention, trọng tâm chủ yếu là nút thắt HBM-SRAM, và trước đây nút thắt này thực sự quan trọng
- Khi HBM nhanh hơn và tensor core tăng tốc nhanh hơn các phần khác của chip, ngay cả độ trễ nhỏ của shared memory cũng trở thành thứ cần loại bỏ hoặc che giấu
- Shared memory được chia thành 32 bank, nên nếu không cẩn thận sẽ xảy ra bank conflict
- Nếu đồng thời yêu cầu nhiều mảnh bộ nhớ khác nhau trong cùng một memory bank, các yêu cầu sẽ bị tuần tự hóa
- Theo kinh nghiệm, kernel có thể chậm đi một cách mất cân đối
- Register layout mà các lệnh WGMMA và MMA yêu cầu có thể gặp bank conflict nếu viết một cách đơn giản
- Giải pháp là tái bố trí shared memory bằng nhiều mẫu swizzling khác nhau để tránh conflict
- Nếu có thể, nên tránh di chuyển dữ liệu giữa register và shared memory; khi cần, tốt hơn là dùng phần cứng dựng sẵn như WGMMA và TMA để di chuyển dữ liệu bất đồng bộ
- Di chuyển đồng bộ bằng warp thực là cách phổ biến nhất, nhưng gần như là fallback tệ nhất
Sinh địa chỉ và TMA
- H100 có tensor core và bộ nhớ đều nhanh, nên bản thân việc sinh memory address cần fetch cũng chiếm một phần đáng kể tài nguyên chip
- Điều này càng rõ hơn khi có thêm các pattern interleaved phức tạp hoặc swizzling pattern
- Tensor Memory Accelerator(TMA) của NVIDIA cho phép chỉ định layout tensor đa chiều của global/shared memory, fetch bất đồng bộ subtile của tensor đó, rồi kích hoạt barrier khi hoàn tất
- TMA giảm chi phí sinh địa chỉ và cũng giúp cấu thành pipeline dễ hơn
- TMA được đánh giá là thiết yếu để khai thác tiềm năng của H100, tương tự
wgmma.mma_async- Theo kinh nghiệm, nó thậm chí có thể quan trọng hơn WGMMA
- Nó tiết kiệm tài nguyên register và instruction dispatch
- Nó cũng có chức năng thực hiện reduction bất đồng bộ vào global memory, hữu ích trong các backward kernel phức tạp
- TMA cũng cần một phần reverse engineering để hiểu swizzling mode, nhưng đỡ đau đớn hơn WGMMA
Chi phí được occupancy che giấu
- Trong CUDA, occupancy chỉ số thread được co-schedule trên cùng phần cứng thực thi
- Warp scheduler của SM quadrant cố issue instruction cho warp đã sẵn sàng nhận lệnh ở mỗi cycle
- H100 có phần ít phụ thuộc vào occupancy hơn các thế hệ trước
- Nhờ các tính năng bất đồng bộ, ngay cả một instruction stream đơn lẻ cũng có thể khiến memory fetch, matrix multiply, shared memory reduction và register math cùng bận rộn
- Tuy vậy, occupancy rất hữu ích để che giấu lỗi và chi phí đồng bộ hóa
- Một pipeline được thiết kế hoàn hảo có thể nhanh mà không cần thêm occupancy
- Quan sát thực tế cho thấy GPU NVIDIA có vẻ được thiết kế với occupancy trong tâm trí
- Vì có nhiều khả năng xảy ra synchronization và sai sót, tăng occupancy thường cải thiện mức sử dụng phần cứng đạt được trong thực tế
- Trên H100, occupancy hữu ích ở mức độ đáng kể, nhưng trên A100 và RTX 4090 thì được xem là còn quan trọng hơn
- Bài viết nhắc tới khả năng nguyên nhân là chúng phụ thuộc nhiều hơn vào synchronous instruction dispatch so với H100
ThunderKittens: DSL nhỏ trong CUDA
- ThunderKittens là DSL nhúng trong CUDA được tạo ra để dễ viết kernel nhanh trên H100
- Ban đầu nó được tạo để dùng nội bộ trong phòng nghiên cứu, rồi sau đó được công khai
- Tên này được đặt vì kittens dễ thương và họ thấy thú vị khi phải gõ
kittens::trong code - ThunderKittens nhắm tới sự đơn giản và cung cấp bốn templated type
- Register tiles: tensor 2D trên register file
- Register vectors: tensor 1D trên register file
- Shared tiles: tensor 2D trong shared memory
- Shared vectors: tensor 1D trong shared memory
- Tile được parameterized theo height, width và layout
- Register vector được parameterized theo length và layout, còn shared vector chỉ dùng length
- shared vector thường không gặp bank conflict
- Các phép toán được cung cấp thao tác tile·vector ở cấp warp hoặc cấp warp group cộng tác
- initializer: chẳng hạn thao tác đặt shared vector về zero
- unary op: như
exp - binary op: như
mul - row/column op: như
row_sum
- ThunderKittens được nhúng trong CUDA, nên không giống các thư viện như Triton, phần trừu tượng của nó “gracefully” thất bại
- Nếu thiếu tính năng, có thể mở rộng theo cách mong muốn
Ví dụ FlashAttention và hiệu năng
- Một kernel forward FlashAttention đơn giản cho RTX 4090 được đưa ra làm ví dụ cho ThunderKittens
- Chỉ xử lý headdim=64
nphải là bội số của 256- Được viết bằng khoảng 60 dòng code CUDA
- Mức sử dụng phần cứng là 75%
- Phần lớn độ phức tạp nằm ở chính thuật toán, không phải ở swizzling pattern hay register layout
- Forward pass FlashAttention-2 cho H100 cũng được viết bằng ThunderKittens
- ThunderKittens bao bọc sự phức tạp của TMA, WGMMA, swizzling mode và descriptor
- Kernel dài khoảng 100 dòng
- Trên H100, nhanh hơn FlashAttention-2 khoảng 30%
- ThunderKittens bao bọc layout và instruction, đồng thời cung cấp primitive, giống một “mini-pytorch” có thể dùng trên GPU
- Based linear attention và các kernel cho những architecture khác sẽ được công khai trong tương lai cũng được phát hành cùng nhau
- Kernel Based linear attention chạy ở 215 TFLOPs
- Nếu xét phần recompute của chính thuật toán, con số vượt 300 TFLOPs
- Linear attention về lý thuyết hiệu quả hơn, nhưng trong lịch sử, trên phần cứng thực tế hiệu quả lại thấp hơn nhiều
- Bài viết cho rằng kết quả này có thể mở rộng phạm vi ứng dụng throughput cao
Tư duy lấy tile làm trung tâm
- Lý do ThunderKittens hoạt động tốt được cho là vì nó không cố làm mọi thứ
- CUDA có sức biểu đạt lớn hơn ThunderKittens rất nhiều
- ThunderKittens là một DSL nhỏ và đơn giản
- Trừu tượng cốt lõi là small tile, và điều này được xem là phù hợp với hướng đi của AI và phần cứng
- ThunderKittens không hỗ trợ chiều nhỏ hơn 16
- Phần cứng cũng được cho là không đặc biệt muốn những chiều nhỏ như vậy
- Bài viết đặt vấn đề kiểu như “nếu phép nhân ma trận nhỏ hơn 16x16 thì có chắc đó là AI không”
- Quan điểm thời CPU xem word 32-bit là register được cho là không phù hợp với phần cứng AI
- Vector register 1024-bit của CUDA được xem là một bước đúng hướng
- Ở đây register là dữ liệu của tile 16x16
- AI vẫn xoay quanh matrix multiply, reduction và reshape, nên trừu tượng tile được xem là phù hợp với cả AI lẫn phần cứng
- Về sau, cần sắp xếp lại các ý tưởng AI theo cách ánh xạ tốt lên phần cứng
- Kích thước recurrent state phải đủ lớn để nằm trong SM
- Compute density không được thấp hơn mức phần cứng yêu cầu
- Điều chỉnh những gì học được từ phần cứng vào thiết kế AI là một hướng quan trọng trong tương lai
Kế hoạch hỗ trợ AMD
- Hỗ trợ AMD hardware cho ThunderKittens sẽ sớm ra mắt
1 bình luận
Ý kiến trên Hacker News
Câu hỏi "Nếu phép nhân ma trận nhỏ hơn 16x16, bạn có chắc đó thực sự là AI không?" khá thú vị
Các yêu cầu đối với phần cứng AI đang ngày càng rõ nét. GPU ban đầu được thiết kế cho mục đích hoàn toàn khác, nhưng được dùng cho AI vì phần cứng nhân ma trận tốt; còn "AI GPU" có thể lược bỏ một số chức năng có trong GPU thực thụ
Cách biểu diễn số cũng đang có xu hướng ngắn hơn, như số thực dấu phẩy động 16-bit, 8-bit, 2-bit, 1-bit, và một ngày nào đó điểm tối ưu sẽ được xác định. Bài viết này cho thấy phần cứng ưu tiên ô 16x16 là khá hợp lý. Rất có thể hiện đã có ai đó đang viết thứ như vậy bằng VHDL, hoặc sắp làm như thế
Cuối cùng có lẽ sẽ xuất hiện những thiết bị đơn giản hơn, ít đa dụng hơn và rẻ hơn, chỉ thực hiện tối đa các phép tính "AI" mà không phải gánh thêm phần cứng không cần thiết
Nvidia có lẽ cũng đang làm việc này, nhưng xét về kinh doanh, việc duy trì một thiết bị gom chung game/giải trí/tiền mã hóa/AI, tức dưới dạng card đồ họa, có thể là lựa chọn tốt hơn
[1] https://github.com/hollance/neural-engine/blob/master/docs/a...
Điều này gợi nhớ đến thời Nervana của Naveen Rao tạo ra driver Nvidia Maxwell nhanh hơn cả driver của chính Nvidia. Không phải mọi lỗi tài liệu của một sản phẩm tăng trưởng nhanh đều là biện pháp đối phó cạnh tranh, nhưng việc các nhà nghiên cứu mất nhiều thời gian để đảo ngược kỹ thuật wgmma, cộng với bối cảnh chính trị Mỹ-Trung quanh H100, khiến Nvidia trông như đang dùng lại chiêu cũ để bảo vệ hào lũy của mình
Vì vậy, thay vì đào quá sâu vào tính đặc thù của H100, nên xem rằng câu hỏi "AI muốn phần cứng như thế nào" cũng bao gồm cả hoàn cảnh thương mại
https://www.amd.com/en/products/accelerators/alveo/v80.html
XDNA Architecture
https://www.amd.com/en/technologies/xdna.html
Đoạn "Lời nói dối của NVIDIA. Đây là cách diễn đạt cực kỳ gây hiểu nhầm về bố cục wgmma swizzled 128b thực tế. Vì sơ đồ này mà tôi mất 3 tuần không thể lấy lại của đời mình, nên đem ra bêu riếu công khai" rất ấn tượng
Không biết liệu có ai ngạc nhiên khi biết rằng một phần khổng lồ trong tiến bộ của AI nằm ở các công việc kỹ thuật như tối ưu hóa phép nhân ma trận, và phần lớn công việc kỹ thuật đó lại là đảo ngược kỹ thuật chip NVIDIA hay không
Bộ lập lịch warp, 4 góc phần tư, bộ tăng tốc bộ nhớ tensor, bố cục wgmma unswizzled…
Ranh giới giữa thuật ngữ GPU và technobabble kiểu Star Trek ngày càng mờ đi
Khi xem các bài khác đôi lúc tôi cũng nghĩ như vậy. Tự hỏi nếu ai đó nhận được link bài ở đây rồi đọc thì sẽ cảm thấy thế nào. Chắc giống như bước vào một buổi tụ họp của dân mê Trek đang bàn về lõi warp
Có lẽ cách tốt nhất để giảm mức tiêu thụ điện năng và tăng tốc độ cho suy luận AI là chuyển sang mạch xấp xỉ analog
Không cần phép nhân và phép cộng số thực dấu phẩy động hoàn hảo; chỉ cần một thiết bị nhận hai điện áp đầu vào và đưa ra điện áp đầu ra đủ gần với kết quả phép nhân
Ưu điểm lớn là thay vì biểu diễn float16 bằng 16 đường dây, ta biểu diễn con số đó bằng điện áp trên 1 đường dây. Về lý thuyết, thậm chí có thể đạt độ chính xác cao hơn nhiều so với float32. Ngoài ra, vì có thể kết nối trực tiếp mà không cần nạp giá trị vào đơn vị số học-logic, mức tiết kiệm diện tích die và điện năng có tiềm năng đạt tới nhiều bậc độ lớn
Ví dụ như chấp nhận việc một trong một triệu bit đầu ra bị lật để cải thiện tỷ lệ hiệu năng/điện năng. Điều này có lẽ khó với float32, nơi chỉ một giá trị vô cực đơn lẻ cũng có thể phá hỏng toàn bộ, nhưng với int8 thì có vẻ vẫn chịu được việc ta muốn 0 mà thỉnh thoảng lại nhận 128
[1] Tôi không chắc đơn vị dấu phẩy động ma trận của H100 trên thực tế có tuân thủ IEEE 754 hay không
Mạng nơ-ron sinh học không gần như kết nối đầy đủ như mạng nơ-ron nhân tạo thông thường; số hệ số kết nối đầu vào/đầu ra của một nơ-ron dưới 10 nên rất cục bộ. Theo những gì ta biết, sinh học cũng không có lan truyền ngược; thay vào đó có phản hồi và vòng lặp
Cũng có thể có các tế bào phụ trợ hoặc quá trình mà ta chưa biết nhưng lại thiết yếu cho chức năng của hệ thần kinh trung ương. Ngay cả ở mức cao cũng có khả năng tồn tại một lượng đáng kể các kết nối được "hard-code", và một số đã được biết đến. Chẳng hạn, các nơ-ron thính giác trong tai được kết nối với nhau và có điều gì đó tương tự tích chập diễn ra để định vị âm thanh. Đây không phải là hiện tượng nổi lên, mà là chức năng có thể có mà không cần huấn luyện
Sự sống đã tìm ra điều này qua hàng tỷ năm và số thế hệ tương tự, nên cũng không đáng ngạc nhiên. Về lý thuyết có thể làm bằng phần mềm, nhưng xét tới hơn 1 nghìn tỷ nơ-ron trong não linh trưởng/người, điều đó cực kỳ khó ngay cả với các máy cỡ nghìn lõi ngày nay. Dù là "cloud" thì cũng sẽ không đáp ứng được độ kết nối và độ trễ cần thiết
Nếu với cách tiếp cận này có thể mô hình hóa thành công cỡ giun hoặc côn trùng thì sẽ rất tuyệt
Bài này làm tôi nhớ lại niềm vui từng cảm nhận trong lớp lập trình song song CS 149
Văn phong của bài này thật sự ấn tượng, và tôi mong được thấy nó trên AMD MI300x. Nếu muốn dùng thời gian trên máy của tôi thì cứ cho biết
Tôi cũng tò mò thực tế nó chạy tốt đến mức nào, hoặc liệu có nên tiết kiệm thêm một chút để mua XTX thay vì 7900 XT, và việc VRAM giảm sẽ ảnh hưởng đến tính hữu dụng thực tế ra sao
Độc giả không nên phải lên tận knowyourmeme.com để hiểu các tác giả muốn nói gì. Tôi thậm chí còn không biết tiêu đề này nghĩa là gì, và tôi nghĩ nó đã lệch mục tiêu đến mức đó
Tôi tò mò nếu muốn hiểu trọn vẹn những bài như thế này thì nên bắt đầu từ đâu, và nên đi theo lộ trình nào
Và bạn nên tự viết một CUDA kernel để thực hiện phép nhân vector-ma trận. Dùng pycuda thì có thể tập trung vào kernel, còn phần còn lại viết bằng Python. Hãy nói với ChatGPT rằng bạn muốn tự tạo một triển khai nhân vector 4000 phần tử với ma trận 4000x12000 và nhờ nó hướng dẫn toàn bộ quá trình
Để thuê GPU thì Runpod khá tốt, hiện có từ GPU giá rẻ đến H100. Ban đầu cứ bắt đầu với GPU cấp thấp là được
Tôi đã dành 2 tháng để triển khai và tối ưu kernel nhân ma trận bằng Spiral
Biểu đồ trong GitHub README (https://github.com/HazyResearch/ThunderKittens/blob/main/att...) rối mắt quá. Loại biểu đồ cột lượn sóng này có hợp pháp không vậy? :P
[1]: https://matplotlib.org/stable/gallery/showcase/xkcd.html#sph...
Cái tên ThunderKittens rất hay. Tôi muốn thấy ThunderKittens xử lý lan truyền ngược FlashAttention, vốn khó hơn lan truyền xuôi khoảng một bậc độ lớn
causal: https://github.com/HazyResearch/ThunderKittens/blob/main/exa...
non-causal: https://github.com/HazyResearch/ThunderKittens/blob/main/exa...
Chẳng phải những nghiên cứu kiểu này đã được các nhóm làm NPU ngày nay thực hiện rồi sao? Ví dụ chip Groq có thể đạt hiệu năng như hiện tại vì dùng kiến trúc chuyên cho AI. Ở phía người tiêu dùng, Apple Silicon cũng khá có năng lực
Tôi không phải người trong lĩnh vực này, nhưng có lẽ chỉ dựa vào các bộ xử lý đa dụng giao tiếp qua những đường tương đối chậm thì sẽ có giới hạn. Nghĩ lại thiết kế ở cấp phần cứng, và cuối cùng hạ giá cho thị trường tiêu dùng, có vẻ là chiến lược dài hạn tốt hơn
Với vài trăm đô có thể mua GPU Nvidia, hoặc mua laptop gaming có 4050 6GB VRAM với giá 900 đô, nên khó gọi AI chạy trên CPU là có năng lực
Ở chỗ làm tôi cũng không có GPU nên đã thử chạy trên CPU, nhưng ngoài việc dùng mô hình nhỏ và chờ đợi thì không thực tế. Cuối cùng tôi phải yêu cầu một máy tính có GPU
“Về mặt kỹ thuật là có thể” và “thực sự dễ dùng” là hai chuyện khác nhau. Nvidia dùng thực sự tốt, còn CPU thì khổ sở và bực bội