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 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 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, khả năng hỗ trợ gốc cho BFloat16 và các loại dữ liệu được định lượng khác (ví dụ: Int4, Int8 hoặc Float8) sẽ khác nhau 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 tới 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ề sự không khớp kiểu này trong 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).
  • Mô phỏng độ chính xác thấp: Trên một số thế hệ phần cứng, 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 chúng 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í chuyển đổi 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, nhưng không đảm bảo có thể 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).