Mã lỗi: E2001

Danh mục: Thời gian biên dịch: RHS DataType không được hỗ trợ trên phần cứng

Lỗi này xảy ra khi kiểu dữ liệu được dùng cho toán hạng Phía bên phải (RHS) trong phép nhân ma trận (ví dụ: jax.lax.dot_general, jax.lax.conv, jax.numpy.matmul hoặc toán tử @) không được thế hệ TPU cụ thể đang dùng hỗ trợ theo cách tự nhiên.

Thông báo lỗi mẫu:

INTERNAL: Mosaic failed to compile TPU kernel: Unsupported matmul RHS type on target: 'vector<256x256xi8>'
...

The MLIR operation involved:
%13440 = "tpu.matmul"(%13435, %13437, %13439) <dimension_numbers = #tpu.dot_dimension_numbers<...>

XLA Backends: TPU

Tổng quan

Đơn vị nhân ma trận (MXU) của TPU hỗ trợ các hoạt động Float32 một cách tự nhiên trên tất cả các thế hệ phần cứng.

Tuy nhiên, tính năng hỗ trợ gốc cho BFloat16 và các loại dữ liệu được lượng tử hoá khác (ví dụ: Int4, Int8 hoặc Float8) tuỳ theo thế hệ phần cứng. Lỗi này xảy ra khi nhân cố gắng ánh xạ một phép nhân ma trận đến MXU bằng cách sử dụng một kiểu dữ liệu mà thế hệ TPU cụ thể của bạn không có mạch điện vật lý để thực thi.

Lỗi này thường cho biết rằng quá trình Chuẩn hoá của trình biên dịch (cố gắng tự động chuyển đổi các loại không được hỗ trợ thành các loại được hỗ trợ (ví dụ: thông qua hoạt động mô phỏng phần mềm)) không tìm thấy quy tắc chuyển đổi hợp lệ hoặc không thể thực hiện việc này vì Chế độ tương thích đã bị tắt.

Gỡ lỗi

Để giải quyết lỗi này, bạn phải điều chỉnh các kiểu dữ liệu cho phù hợp với khả năng của phần cứng. Bạn có thể chọn một trong các phương án sau đây:

1. Truyền đến các loại gốc

Cách khắc phục đáng tin cậy nhất là truyền các toán hạng theo cách thủ công sang kiểu dữ liệu được phần cứng hỗ trợ (chẳng hạn như Float32 hoặc BFloat16 trên TPU phiên bản 4 trở lên) bên trong nhân trước khi thực hiện thao tác matmul.

  • Lý do: Float32 là loại dữ liệu phổ biến được MXU hỗ trợ nguyên bản trên mọi thế hệ TPU.
  • Đánh đổi: Điều này đi kèm với chi phí VPU (Đơn vị xử lý vectơ) – số chu kỳ cần thiết để thực hiện truyền, nhưng đảm bảo rằng hạt nhân của bạn sẽ chạy trên phần cứng hiện tại.

2. Kiểm tra Chế độ tương thích

Thông thường, trình biên dịch có thể tự động xử lý các vấn đề về kiểu không khớp này ở Chế độ tương thích (được bật theo mặc định). Kiểm tra kỹ các cấu hình XLA để đảm bảo --xla_mosaic_compat_mode không được đặt thành false.

Điều này đóng vai trò là một "polyfill", chèn các chuỗi mô phỏng phần mềm cho các thao tác mà phần cứng của bạn không hỗ trợ theo cách tự nhiên.

Những tính năng mà Chế độ tương thích cho phép:

  • MatMul có độ chính xác hỗn hợp: Cho phép kết hợp các toán hạng Integer với các bộ tích luỹ Float bằng cách tự động chèn các thao tác truyền (ví dụ: mở rộng số nguyên thành Float32 trước matmul).
  • Hoạt động mô phỏng có độ chính xác thấp: Trên một số thế hệ phần cứng nhất định, hoạt động này mô phỏng các loại không được hỗ trợ như dấu phẩy động 4-bit (4E2M1FN) hoặc dấu phẩy động 8-bit (8E4M3FN) bằng cách mở rộng các loại này thành các loại được hỗ trợ như BFloat16 hoặc Float32 trước khi thực thi.

Xin lưu ý rằng chế độ này ưu tiên khả năng tương thích hơn hiệu suất cao nhất vì hoạt động mô phỏng yêu cầu thêm các chỉ dẫn để chuyển đổi định dạng dữ liệu trước khi MXU có thể hoạt động trên các định dạng đó.

3. Nâng cấp phần cứng hoặc yêu cầu hỗ trợ

Nếu thuật toán của bạn hoàn toàn yêu cầu hiệu suất gốc cho các loại như Int4 hoặc Float8 mà không cần chi phí truyền hoặc mô phỏng, thì bạn sẽ cần chạy trên TPU thế hệ mới hơn có hỗ trợ gốc.

Yêu cầu về tính năng: Nếu bạn cho rằng phần cứng của mình hỗ trợ thao tác này hoặc nếu trình biên dịch thiếu một đường dẫn mô phỏng hợp lệ ngay cả ở Chế độ tương thích, vui lòng gửi yêu cầu về tính năng. Chúng tôi thường đảm bảo rằng các thao tác tương thích với phiên bản sau. Vì vậy, nếu nhân của bạn chạy trên một thế hệ TPU, thì nhân đó sẽ chạy trên tất cả các thế hệ trong tương lai. Tuy nhiên, không có gì đảm bảo rằng sẽ có tính năng mô phỏng cho các thế hệ cũ (đối với một số thế hệ, việc truyền sẽ rất tốn kém).