Các tính năng nâng cao WebAssembly và WebGPU để AI trên web nhanh hơn, phần\"2

Tài liệu này là nội dung tiếp nối các điểm cải tiến về WebAssembly và WebGPU để AI trên web nhanh hơn, phần 1. Bạn nên đọc bài đăng này hoặc xem bài nói chuyện tại IO 24 trước khi tiếp tục.

Austin Eng
Austin Eng
Deepti Gandluri
Deepti Gandluri
François Beaufort
François Beaufort

WebGPU

WebGPU cung cấp cho các ứng dụng web quyền truy cập vào phần cứng GPU của máy khách để thực hiện việc tính toán hiệu quả, song song cao độ. Kể từ khi ra mắt WebGPU trong Chrome, chúng tôi đã được chứng kiến những bản minh hoạ tuyệt vời về trí tuệ nhân tạo (AI) và học máy (ML) trên web.

Ví dụ: công cụ Web Stable Diffusion (Khuếch tán ổn định cho web) đã chứng minh rằng trình duyệt có thể sử dụng AI để tạo hình ảnh từ văn bản. Đầu năm nay, nhóm Mediapipe của Google đã công bố hỗ trợ thử nghiệm cho suy luận mô hình ngôn ngữ lớn.

Ảnh động sau đây minh hoạ Gemma, mô hình ngôn ngữ lớn (LLM) nguồn mở của Google, chạy hoàn toàn trên thiết bị trong Chrome theo thời gian thực.

Bản minh hoạ của ôm khuôn mặt sau đây về Mô hình Segment Anything (Phân khúc bất cứ thứ gì) của Meta tạo ra các mặt nạ đối tượng chất lượng cao hoàn toàn trên máy khách.

Đây chỉ là một vài dự án tuyệt vời thể hiện sức mạnh của WebGPU dành cho trí tuệ nhân tạo và công nghệ học máy. WebGPU cho phép các mô hình này và các mô hình khác chạy nhanh hơn đáng kể so với tốc độ trên CPU.

Điểm chuẩn WebGPU dành cho tính năng nhúng văn bản của Ôm Face cho thấy tốc độ tăng rất nhanh so với việc triển khai CPU của cùng một mô hình. Trên máy tính xách tay Apple M1 Max, WebGPU nhanh hơn 30 lần. Những người khác báo cáo rằng WebGPU tăng tốc độ đo điểm chuẩn lên hơn 120 lần.

Cải thiện các tính năng của WebGPU cho trí tuệ nhân tạo và công nghệ học máy

WebGPU rất phù hợp với các mô hình trí tuệ nhân tạo (AI) và học máy (ML) vốn có thể chứa hàng tỷ tham số nhờ tính năng hỗ trợ chương trình đổ bóng điện toán. Chương trình đổ bóng điện toán chạy trên GPU và giúp chạy các thao tác mảng song song trên khối lượng lớn dữ liệu.

Bên cạnh nhiều điểm cải tiến đối với WebGPU trong năm qua, chúng tôi đã tiếp tục bổ sung thêm nhiều chức năng để cải thiện hiệu suất của công nghệ học máy và trí tuệ nhân tạo (AI) trên web. Gần đây, chúng tôi đã ra mắt hai tính năng mới: dấu phẩy động 16 bit và các sản phẩm dấu chấm số nguyên đóng gói.

Dấu phẩy động 16 bit

Hãy nhớ rằng khối lượng công việc học máy không đòi hỏi độ chính xác. shader-f16 là một tính năng cho phép sử dụng loại f16 trong ngôn ngữ tô bóng WebGPU. Loại dấu phẩy động này chiếm 16 bit, thay vì 32 bit thông thường. f16 có phạm vi nhỏ hơn và kém chính xác hơn, nhưng đối với nhiều mô hình học máy, điều này là đủ.

Tính năng này giúp tăng hiệu quả theo một số cách:

  • Giảm bộ nhớ: Tensor có các phần tử f16 chiếm một nửa dung lượng, giúp giảm một nửa mức sử dụng bộ nhớ. Các phép tính GPU thường gây tắc nghẽn trên băng thông bộ nhớ, vì vậy, một nửa bộ nhớ thường có nghĩa là chương trình đổ bóng chạy nhanh gấp đôi. Về mặt kỹ thuật, bạn không cần F16 để tiết kiệm băng thông bộ nhớ. Có thể lưu trữ dữ liệu ở định dạng có độ chính xác thấp, sau đó mở rộng đến f32 đầy đủ trong chương trình đổ bóng để tính toán. Tuy nhiên, GPU cần thêm công suất tính toán để đóng gói và giải nén dữ liệu.

  • Giảm khả năng chuyển đổi dữ liệu: f16 sử dụng ít điện toán hơn bằng cách giảm thiểu số lượt chuyển đổi dữ liệu. Dữ liệu có độ chính xác thấp có thể được lưu trữ rồi sử dụng trực tiếp mà không cần chuyển đổi.

  • Tăng cường tính song song: Các GPU hiện đại có thể điều chỉnh đồng thời nhiều giá trị hơn trong đơn vị thực thi của GPU, cho phép thực hiện nhiều phép tính song song hơn. Chẳng hạn, một GPU hỗ trợ lên đến 5 nghìn tỷ f32 hoạt động dấu phẩy động mỗi giây có thể hỗ trợ 10 nghìn tỷ f16 hoạt động dấu phẩy động mỗi giây.

Ảnh chụp màn hình điểm chuẩn WebGPU cho tính năng nhúng văn bản
Với shader-f16, điểm chuẩn WebGPU dành cho nhúng văn bản của Hugging Face có thể chạy điểm chuẩn nhanh hơn 3 lần so với f32 trên máy tính xách tay Apple M1 Max.

WebLLM là một dự án có thể chạy nhiều mô hình ngôn ngữ lớn. Lớp này sử dụng Apache TVM, một khung trình biên dịch học máy nguồn mở.

Tôi đã yêu cầu WebLLM lên kế hoạch cho một chuyến đi đến Paris bằng cách sử dụng mô hình thông số Llama 3 8 tỷ. Kết quả cho thấy trong giai đoạn điền sẵn của mô hình, f16 nhanh hơn 2,1 lần so với f32. Trong giai đoạn giải mã, tốc độ này nhanh hơn 1,3 lần.

Trước tiên, các ứng dụng phải xác nhận rằng bộ điều hợp GPU hỗ trợ f16 và nếu bộ điều hợp GPU có sẵn, hãy bật một cách rõ ràng khi yêu cầu thiết bị GPU. Nếu f16 không được hỗ trợ, bạn không thể yêu cầu nó trong mảng requiredFeatures.

// main.js

const adapter = await navigator.gpu.requestAdapter();
const supportsF16 = adapter.features.has('shader-f16');
if (supportsF16) {
  // Use f16.
  const device = await adapter.requestDevice({
    requiredFeatures: ['shader-f16'],
  });
  initApp(device);
}

Sau đó, trong chương trình đổ bóng WebGPU, bạn phải bật f16 một cách rõ ràng ở trên cùng. Sau đó, bạn có thể sử dụng dữ liệu này trong chương trình đổ bóng như mọi kiểu dữ liệu số thực khác.

// my-shader.wgsl

enable f16;

struct Data {
  values : array<vec4<f16>>
}
@group(0) @binding(0) var<storage, read> data : Data;
@compute @workgroup_size(64) fn main(@builtin(global_invocation_id) gid : vec3u) {
  let value : vec4<f16> = data.values[gid.x];
  ...
}

Sản phẩm dấu chấm số nguyên được đóng gói

Nhiều kiểu máy vẫn hoạt động tốt với chỉ 8 bit độ chính xác (một nửa khẩu độ f16). Phương pháp này phổ biến trong các LLM và mô hình hình ảnh để phân đoạn và nhận dạng đối tượng. Tuy nhiên, chất lượng đầu ra cho các mô hình bị suy giảm với độ chính xác thấp hơn, vì vậy việc lượng tử hoá 8 bit không phù hợp với mọi ứng dụng.

Tương đối ít GPU hỗ trợ sẵn các giá trị 8 bit. Đây là nơi các sản phẩm dấu chấm số nguyên đóng gói xuất hiện. Chúng tôi đã vận chuyển DP4a trong Chrome 123.

Các GPU hiện đại có hướng dẫn đặc biệt để lấy hai số nguyên 32 bit, diễn giải chúng là 4 số nguyên 8 bit được đóng gói liên tiếp và tính toán tích vô hướng giữa các thành phần của chúng.

Điều này đặc biệt hữu ích cho AI và công nghệ học máy vì hạt nhân nhân ma trận bao gồm rất nhiều sản phẩm chấm.

Ví dụ: hãy nhân ma trận 4 x 8 với vectơ 8 x 1. Quá trình tính toán bao gồm việc lấy các sản phẩm có 4 dấu chấm để tính toán từng giá trị trong vectơ đầu ra; A, B, C và D.

Ví dụ về sơ đồ nhân vectơ của ma trận

Quy trình tính toán từng dữ liệu đầu ra này đều giống nhau; chúng ta sẽ xem các bước để tính toán một trong số các bước đó. Trước khi thực hiện một phép tính nào đó, đầu tiên chúng ta cần chuyển đổi dữ liệu số nguyên 8 bit thành kiểu dữ liệu có thể dùng để thực hiện phép tính, chẳng hạn như f16. Sau đó, chúng ta chạy phép nhân theo từng phần tử và cuối cùng là cộng tất cả các sản phẩm lại với nhau. Tổng cộng, đối với toàn bộ phép nhân vectơ ma trận, chúng ta thực hiện 40 phép nhân số nguyên để giải nén dữ liệu, 32 phép nhân số thực có độ chính xác đơn và 28 phép cộng số thực có độ chính xác đơn.

Đối với ma trận lớn hơn có nhiều thao tác hơn, các sản phẩm dấu chấm số nguyên được đóng gói có thể giúp giảm khối lượng công việc.

Đối với mỗi dữ liệu đầu ra trong vectơ kết quả, chúng ta thực hiện hai thao tác với sản phẩm dạng dấu chấm đóng gói bằng Ngôn ngữ tạo bóng WebGPU tích hợp dot4U8Packed, sau đó cộng các kết quả lại với nhau. Tổng cộng, đối với toàn bộ phép nhân vectơ của ma trận, chúng tôi sẽ không thực hiện bất kỳ chuyển đổi dữ liệu nào. Chúng ta thực thi các sản phẩm có 8 dấu chấm và bổ sung 4 số nguyên.

Sơ đồ ví dụ về phép nhân ma trận số nguyên đóng gói

Chúng tôi đã thử nghiệm các sản phẩm dấu chấm số nguyên đóng gói với dữ liệu 8 bit trên nhiều GPU tiêu dùng. So với dấu phẩy động 16 bit, chúng ta có thể thấy rằng 8 bit nhanh hơn từ 1,6 đến 2,8 lần. Khi chúng tôi sử dụng thêm các sản phẩm dấu chấm số nguyên đóng gói, hiệu suất còn tốt hơn nữa. Nhanh hơn 1,7 đến 2,9 lần.

Ảnh chụp màn hình tăng tốc nhân lên vectơ của ma trận: f16 so với u8
Biểu đồ 1: Tăng tốc vectơ ma trận, so sánh f16 với U8 và U8 với chấm4U8Packed.

Kiểm tra xem thuộc tính wgslLanguageFeatures có hỗ trợ trình duyệt hay không. Nếu GPU vốn không hỗ trợ các sản phẩm chấm đóng gói sẵn, thì trình duyệt sẽ tự thực hiện việc triển khai polyfill.

// main.js

if (navigator.gpu.wgslLanguageFeatures.has('packed_4x8_integer_dot_product')) {
  // Use dot4U8Packed, dot4I8Packed builtin
  // functions in the shaders.
}

Đoạn mã sau đây phân biệt (khác biệt) làm nổi bật những thay đổi cần thiết để sử dụng sản phẩm số nguyên được đóng gói trong chương trình đổ bóng WebGPU.

Trước – Chương trình đổ bóng WebGPU tích luỹ một phần các sản phẩm là dấu chấm vào biến "tổng". Ở cuối vòng lặp, `sum` lưu giữ tích dấu chấm đầy đủ giữa một vectơ và một hàng của ma trận đầu vào.

// my-dot-product.wgsl

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid : vec3u) {
  var sum : f16;
  let start = gid.x * uniforms.dim;
  for (var i = 0u; i < uniforms.dim; i++) {
    let v1 : vec4<f16> = vector.values[i];
    let v2 : vec4<f16> = matrix.values[start + i];
    sum += dot(v1, v2);
  }
}

Sau – Chương trình đổ bóng WebGPU được viết để sử dụng các sản phẩm dấu chấm số nguyên được đóng gói. Sự khác biệt chính là thay vì tải 4 giá trị float ra khỏi vectơ và ma trận, chương trình đổ bóng này tải một số nguyên 32 bit duy nhất. Số nguyên 32 bit này chứa dữ liệu của bốn giá trị số nguyên 8 bit. Sau đó, chúng ta gọi dot4U8Packed để tính tích vô hướng của hai giá trị.

// my-dot-product.wgsl

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid : vec3u) {
  var sum : f32;
  let start = gid.x * uniforms.dim;
  for (var i = 0u; i < uniforms.dim; i++) {
    let v1 : u32 = vector.values[i];
    let v2 : u32 = matrix.values[start + i];
    sum += dot4U8Packed(v1, v2);
  }
}

Cả sản phẩm dấu phẩy động 16 bit và dấu chấm số nguyên đóng gói đều là những tính năng xuất xưởng trong Chrome, giúp tăng tốc độ AI và công nghệ học máy. Dấu phẩy động 16 bit được cung cấp khi phần cứng hỗ trợ dấu phẩy động 16 bit đó. Đồng thời, Chrome triển khai các sản phẩm dấu chấm số nguyên đóng gói trên tất cả thiết bị.

Bạn có thể dùng các tính năng này trong Bản ổn định của Chrome ngay hôm nay để đạt được hiệu suất tốt hơn.

Các tính năng được đề xuất

Sắp tới, chúng tôi sẽ điều tra hai tính năng khác: nhóm con và ma trận hợp tác nhân lên.

Tính năng nhóm con cho phép tính song song ở cấp SIMD giao tiếp hoặc thực hiện các phép toán tập thể, chẳng hạn như tổng cho hơn 16 số. Điều này giúp bạn chia sẻ dữ liệu trên nhiều luồng một cách hiệu quả. Các nhóm con được hỗ trợ trên các API GPU hiện đại, với nhiều tên và hình thức hơi khác nhau.

Chúng tôi đã tóm tắt tập hợp chung thành một đề xuất mà chúng tôi đưa vào nhóm tiêu chuẩn hoá WebGPU. Đồng thời, chúng tôi đã xây dựng các nhóm con nguyên mẫu trong Chrome dựa trên một cờ thử nghiệm và đưa kết quả ban đầu vào thảo luận. Vấn đề chính là làm thế nào để đảm bảo hoạt động di động.

Ma trận hợp tác nhân lên là tính năng bổ sung mới hơn cho GPU. Bạn có thể chia một phép nhân ma trận lớn thành nhiều phép nhân ma trận nhỏ hơn. Phép nhân ma trận cộng tác thực hiện phép nhân trên các khối có kích thước cố định nhỏ hơn này theo một bước logic duy nhất. Trong bước đó, một nhóm các luồng sẽ cộng tác hiệu quả để tính toán kết quả.

Chúng tôi đã khảo sát khả năng hỗ trợ đối với các API GPU cơ bản và dự định trình bày một đề xuất cho nhóm tiêu chuẩn hoá WebGPU. Tương tự như với các nhóm con, chúng tôi hy vọng rằng phần lớn nội dung thảo luận sẽ tập trung vào khả năng di chuyển.

Để đánh giá hiệu suất hoạt động của nhóm con, trong một ứng dụng thực tế, chúng tôi đã tích hợp hỗ trợ thử nghiệm cho các nhóm con vào MediaPipe và thử nghiệm tính năng này với nguyên mẫu của Chrome cho hoạt động nhóm con.

Chúng tôi đã sử dụng các nhóm con trong hạt nhân GPU của giai đoạn điền sẵn của mô hình ngôn ngữ lớn, vì vậy, tôi chỉ báo cáo tốc độ cho giai đoạn điền sẵn. Trên GPU Intel, chúng ta thấy rằng các nhóm con hoạt động nhanh hơn gấp 2,5 lần so với đường cơ sở. Tuy nhiên, những cải tiến này không nhất quán trên các GPU khác nhau.

Ảnh chụp màn hình tăng tốc nhóm con trong suy luận LLM MediaPipe
Biểu đồ 2 Các nhóm con giúp tính năng điền sẵn chạy nhanh hơn 2,5 lần trên GPU Intel Core Lake GT2, với hỗ trợ thử nghiệm trong Chrome và Mediapipe.

Biểu đồ tiếp theo cho thấy kết quả của việc áp dụng các nhóm con để tối ưu hoá ma trận nhân microbenchmark trên nhiều GPU của người dùng. Nhân ma trận là một trong những thao tác nặng hơn trong các mô hình ngôn ngữ lớn. Dữ liệu cho thấy trên nhiều GPU, các nhóm con tăng tốc độ 2, 5 và thậm chí 13 lần so với tốc độ cơ sở. Tuy nhiên, lưu ý rằng trên GPU đầu tiên, các nhóm con cũng không tốt hơn nhiều.

Ảnh chụp màn hình về tính năng tăng tốc nhóm con cho phép nhân ma trận
Biểu đồ 3 Việc áp dụng các nhóm con để nhân ma trận có thể làm tăng hiệu suất hơn nữa.

Gặp khó khăn trong việc tối ưu hoá GPU

Suy cho cùng, cách tốt nhất để tối ưu hoá GPU của bạn phụ thuộc vào GPU mà ứng dụng cung cấp. Việc sử dụng các tính năng GPU mới lạ mắt không phải lúc nào cũng mang lại hiệu quả như bạn mong đợi vì có thể có nhiều yếu tố phức tạp liên quan. Chiến lược tối ưu hoá tốt nhất trên một GPU chưa chắc là chiến lược tốt nhất trên một GPU khác.

Bạn muốn giảm thiểu băng thông bộ nhớ trong khi vẫn sử dụng hoàn toàn các luồng điện toán của GPU.

Các mẫu truy cập bộ nhớ cũng có thể thực sự quan trọng. GPU có xu hướng hoạt động tốt hơn nhiều khi các luồng điện toán truy cập vào bộ nhớ theo mẫu tối ưu cho phần cứng. Lưu ý quan trọng: Bạn sẽ thấy có đặc điểm hiệu suất khác nhau trên mỗi phần cứng GPU. Có thể bạn cần chạy các biện pháp tối ưu hoá khác nhau tuỳ thuộc vào GPU.

Trong biểu đồ sau, chúng tôi đã sử dụng cùng một thuật toán nhân ma trận, nhưng thêm một phương diện khác để minh hoạ rõ hơn tác động của các chiến lược tối ưu hoá khác nhau, cũng như sự phức tạp và phương sai giữa các GPU khác nhau. Tại đây, chúng tôi đã giới thiệu một kỹ thuật mới, gọi là "Swizzle". Swizzle tối ưu hoá các mẫu truy cập bộ nhớ sao cho phần cứng tối ưu hơn.

Bạn có thể thấy tình trạng tăng tốc bộ nhớ có tác động đáng kể; đôi khi lại có tác động mạnh mẽ hơn so với các nhóm con. Trên GPU 6, swizzle cung cấp tốc độ tăng 12 lần, trong khi các nhóm con cung cấp tốc độ tăng 13 lần. Khi kết hợp, chúng có tốc độ tăng tốc đáng kinh ngạc gấp 26 lần. Đối với các GPU khác, đôi khi swizzle và các nhóm con kết hợp lại hoạt động tốt hơn so với một nhóm riêng lẻ. Còn trên các GPU khác, chỉ sử dụng swizzle sẽ mang lại hiệu quả cao nhất.

Ảnh chụp màn hình tính năng tăng tốc cho chiến lược nhân ma trận
Biểu đồ 4.

Việc tinh chỉnh và tối ưu hoá thuật toán GPU để hoạt động tốt trên mọi phần cứng có thể đòi hỏi nhiều kiến thức chuyên môn. Nhưng may mắn là có rất nhiều tài năng đang phát triển cho các khung thư viện cấp cao hơn, chẳng hạn như Mediapipe, Transformers.js, Apache TVM, ONNX Runtime Web và nhiều hơn nữa.

Các thư viện và khung có lợi thế để xử lý sự phức tạp trong việc quản lý nhiều cấu trúc GPU, cũng như tạo mã dành riêng cho nền tảng có thể chạy tốt trên máy khách.

Cướp lại bóng

Nhóm Chrome tiếp tục giúp phát triển các tiêu chuẩn WebAssembly và WebGPU để cải thiện nền tảng web cho các tải công việc của công nghệ học máy. Chúng tôi đang đầu tư vào việc tính toán các dữ liệu gốc nhanh hơn, khả năng tương tác tốt hơn trên các tiêu chuẩn web và đảm bảo rằng các mô hình cả lớn và nhỏ đều có thể chạy hiệu quả trên nhiều thiết bị.

Mục tiêu của chúng tôi là tối đa hoá các khả năng của nền tảng trong khi vẫn duy trì những điểm tốt nhất trên web: phạm vi tiếp cận, khả năng hữu dụng và tính di động. Và chúng tôi không làm việc này một mình. Chúng tôi đang cộng tác với các nhà cung cấp trình duyệt khác tại W3C và nhiều đối tác phát triển.

Chúng tôi hy vọng bạn còn nhớ những điều sau khi làm việc với WebAssembly và WebGPU:

  • Suy luận bằng AI hiện đã có trên web, trên nhiều thiết bị. Điều này mang lại lợi thế khi chạy trên các thiết bị khách, chẳng hạn như giảm chi phí máy chủ, độ trễ thấp và tăng cường quyền riêng tư.
  • Mặc dù nhiều tính năng được thảo luận chủ yếu có liên quan đến các tác giả khung, các ứng dụng của bạn có thể hưởng lợi mà không mất nhiều chi phí.
  • Các tiêu chuẩn web luôn thay đổi và phát triển không ngừng, và chúng tôi luôn mong muốn nhận được ý kiến phản hồi. Hãy chia sẻ mã của bạn cho WebAssemblyWebGPU.

Xác nhận

Chúng tôi muốn cảm ơn đội đồ hoạ web Intel, những người đã góp phần thúc đẩy WebGPU f16 và đóng gói các tính năng sản phẩm dấu chấm nguyên. Chúng tôi muốn cảm ơn các thành viên khác trong nhóm làm việc WebAssembly và WebGPU tại W3C, bao gồm cả các nhà cung cấp trình duyệt khác.

Xin chân thành cảm ơn các nhóm phụ trách AI và ML (học máy và trí tuệ nhân tạo) tại Google và trong cộng đồng nguồn mở vì đã luôn là những đối tác tuyệt vời. Và tất nhiên là tất cả những người trong nhóm của chúng tôi, những người đã giúp tất cả những điều này có thể thực hiện được.