Đặc điểm của GPU H100
- Cung cấp 80GB bộ nhớ HBM3 và băng thông 3TB/s (thực tế hơi thấp hơn một chút)
- Cung cấp 50MB cache L2 và băng thông 12TB/s. Trên GPU, phần này được chia thành hai vùng 25MB và kết nối bằng crossbar (crossbar là một yếu tố làm suy giảm hiệu năng)
- Gồm 132 Streaming Multiprocessor (SM), mỗi SM có cấu hình như sau:
- Cung cấp tối đa 227KB shared memory trong 256KB cache L1 (tổng cộng khoảng 33TB/s băng thông)
- Cung cấp Tensor Memory Accelerator (TMA) cho phép tạo địa chỉ bất đồng bộ và fetching bộ nhớ. Ngoài ra còn hỗ trợ các tính năng như mạng bộ nhớ on-chip, nhưng bài viết này không đề cập
- Được chia thành 4 quadrant, mỗi quadrant gồm warp scheduler, 512 vector register (mỗi register chứa 32 word 4-byte), tensor core cho phép nhân ma trận, và các instruction tích hợp hỗ trợ tính toán song song như sum/multiply
Mẹo tối ưu hiệu năng GPU H100
- Điều quan trọng là phải tận dụng Tensor Core tối đa. H100 có hiệu năng nhân ma trận FP16 đạt 989 TFLOPs, nên mức độ tận dụng Tensor Core ảnh hưởng rất lớn đến tỷ lệ sử dụng toàn bộ GPU
- Tuy nhiên, để tận dụng Tensor Core tối đa, cần cân nhắc các điểm sau:
- Việc sử dụng lệnh Warp Group Matrix Multiply Accumulate (WGMMA) là bắt buộc nhưng khá khó dùng
- Shared Memory không nhanh như tưởng tượng, và cần sử dụng rất cẩn thận
- Việc tạo địa chỉ có chi phí lớn, nên cần tối ưu bằng cách tận dụng Tensor Memory Accelerator (TMA) v.v.
- Tăng occupancy vẫn tiếp tục hữu ích, và register là tài nguyên chủ chốt
- Những đặc điểm này phần nào cũng áp dụng cho các GPU khác chứ không chỉ H100, nhưng việc tận dụng Tensor Core đặc biệt quan trọng và cũng khó hơn trên H100
ThunderKittens: CUDA embedded DSL được tối ưu cho H100
- Một CUDA-based embedded DSL được phát triển để khai thác tối đa hiệu năng của các GPU hiện đại như NVIDIA H100
- Cung cấp 4 kiểu template dựa trên tile như sau:
- Register Tile (tensor 2D được lưu trong register)
- Register Vector (tensor 1D được lưu trong register)
- Shared Tile (tensor 2D được lưu trong Shared Memory)
- Shared Vector (tensor 1D được lưu trong Shared Memory)
- Ngoài ra còn cung cấp nhiều toán tử để thao tác tile (exp, mul, sum, v.v.) ở cấp warp hoặc warp group
- Khi triển khai các kernel Flash Attention và Flash Attention 2 hiện có bằng ThunderKittens, mã nguồn trở nên gọn hơn đáng kể và hiệu năng trên H100 tăng tối đa 30%
- Kernel Based Linear Attention cũng được triển khai bằng ThunderKittens và đạt hiệu năng 215 TFLOPs (nếu tính cả recompute theo đặc tính thuật toán thì vượt 300 TFLOPs)
Suy ngẫm từ góc nhìn triết học
- Lý do ThunderKittens hoạt động tốt là vì nó không cố hỗ trợ mọi thứ, mà thay vào đó cung cấp một abstraction dựa trên tile vừa đơn giản vừa khớp tốt với kiến trúc GPU
- Việc dùng vector register 1024-bit thay cho word 32-bit truyền thống cũng là một bước tiến, nhưng cần một sự chuyển đổi tư duy để coi tile 16x16 là đơn vị của register
- Xét đến việc workload AI rốt cuộc chủ yếu xoay quanh nhân ma trận, reduction và reshape, cách tiếp cận dựa trên tile là hợp lý; từ góc độ phần cứng, ngoài systolic array, phần cứng cũng sẽ tiến hóa theo hướng hỗ trợ các phép nhân ma trận nhỏ
- Xa hơn nữa, cần chuyển đổi cách tư duy sang thiết kế thuật toán AI dưới dạng được tối ưu hóa cho phần cứng. Ví dụ, cần giới hạn kích thước trạng thái của RNN sao cho vừa trong SM, và điều chỉnh mật độ tính toán theo mức mà phần cứng yêu cầu
Ý kiến của GN⁺
- ThunderKittens có thể là một lựa chọn hấp dẫn với các nhà phát triển đã quen với CUDA, vì có thể dễ dàng đạt được cải thiện hiệu năng mà không cần chỉnh sửa quá nhiều mã kernel hiện có
- Tuy nhiên, với người mới bắt đầu thì rào cản gia nhập vẫn có thể khá cao. Có lẽ sẽ cần thêm nhiều ví dụ mã và tài liệu học tập đa dạng hơn trong tương lai
- Kế hoạch mở rộng hỗ trợ ThunderKittens sang các GPU khác như AMD chứ không chỉ H100 cũng rất đáng chú ý. Điều này có thể góp phần giảm sự phụ thuộc vào một vendor duy nhất
- Cuối cùng, việc thiết kế chính mô hình/thuật toán AI dưới dạng tối ưu cho phần cứng là một điểm cực kỳ quan trọng. Để làm được điều đó, cần có sự hiểu biết sâu về đặc tính phần cứng, và ThunderKittens có thể giúp các nhà phát triển có được những insight như vậy
- Kỳ vọng rằng hoạt động R&D liên tục và các đóng góp mã nguồn mở của Hazy Research sẽ giúp hệ sinh thái CUDA sôi động hơn. Tuy vậy, về dài hạn, dường như vẫn cần sự xuất hiện của những framework có mức độ abstraction cao hơn
1 bình luận
Ý kiến trên Hacker News
Các yêu cầu đối với phần cứng AI đang ngày càng rõ ràng hơn. GPU ban đầu được thiết kế cho mục đích khác, nhưng lại được dùng cho AI nhờ có phần cứng nhân ma trận tốt. “AI GPU” có thể lược bỏ một số chức năng của GPU thực thụ, và xu hướng đang hướng tới các định dạng số ngắn hơn (dấu phẩy động 16 bit, 8 bit, 2 bit, 1 bit). Bài báo này gợi ý rằng phần cứng ưa thích các tile 16x16 có nhiều lợi thế.
Cần có bộ đồng xử lý chuyên dụng cho AI (NPU), đặc biệt trong các hệ thống desktop prosumer dành cho lập trình viên, chuyên gia, game thủ và những nhóm tương tự. GPU hoạt động được trong môi trường doanh nghiệp, nhưng khá bất tiện khi dùng cho AI ở góc độ điện toán cá nhân. Đặc biệt, giới hạn VRAM và việc thiếu API mở tiêu chuẩn ngoài Vulkan là vấn đề.
Muốn thúc đẩy nghiên cứu AI thì cần nghiên cứu thần kinh học và tâm lý học tốt hơn. Ngoài ra, những vấn đề liên quan đến topology đồ thị của mạng nơ-ron cũng có thể có liên quan.
Đây có phải là CUTLASS được làm cho thân thiện hơn với người dùng không?
Linh vật ThunderKittens mang cảm giác như mèo/Sony Aibo. Có vẻ được tạo ra khá tốt bằng AI.
Nếu Universal Basic Compute (UBC) được xem như một giải pháp thay thế cho universal basic income, thì đó sẽ là một tương lai rất phản địa đàng. Hãy tưởng tượng một công ty như Nvidia tạo ra toàn bộ năng lực tính toán.