Lượng tử hoá StableHLO

Các loại lượng tử hoá trong StableHLO

Lượng tử hoá là một kỹ thuật tối ưu hoá các mô hình học máy bằng cách chuyển đổi số thực dấu phẩy động (chẳng hạn như số được dùng trong các mô hình ban đầu) thành số nguyên có độ chính xác thấp hơn. Điều này giúp giảm mức sử dụng bộ nhớ và tăng tốc độ tính toán, giúp các mô hình hiệu quả hơn khi triển khai trên các thiết bị có tài nguyên hạn chế.

Lượng tử hoá StableHLO tuân theo quy cách lượng tử hoá LiteRT, sử dụng một lược đồ lượng tử hoá đồng nhất có hỗ trợ cả lượng tử hoá theo trục và theo từng tensor. Loại này kế thừa biểu thức loại từ phương ngữ Quant của MLIR, cung cấp một cách tiêu chuẩn hoá để biểu thị các loại dữ liệu được lượng tử hoá.

Lượng tử hoá đồng nhất ánh xạ các giá trị dấu phẩy động thành số nguyên bằng cách sử dụng kích thước bước đồng nhất, dẫn đến các giá trị được lượng tử hoá có khoảng cách đều nhau. Điều này đạt được thông qua mối quan hệ ái lực bằng cách sử dụng hai tham số định lượng khoá.

Việc lượng tử hoá đồng nhất giúp đơn giản hoá việc biểu diễn các số dấu phẩy động bằng cách ánh xạ các số đó thành các số nguyên có khoảng cách đều nhau. Hoạt động ánh xạ này được thực hiện thông qua một phép biến đổi affine sử dụng hai tham số chính: tỷ lệ và điểm không. Thang đo xác định kích thước bước giữa các giá trị được lượng tử hoá liên tiếp. Thang đo nhỏ hơn có nghĩa là các giá trị được lượng tử hoá gần nhau hơn. Điểm 0 xác định giá trị số nguyên đại diện cho số 0 trong không gian dấu phẩy động ban đầu.

Mối quan hệ giữa giá trị dấu phẩy động ban đầu (real_value) và giá trị số nguyên được lượng tử hoá (quantized_value) trong quá trình lượng tử hoá đồng nhất là:

real_value = scale * (quantized_value - zero_point)

Lượng tử hoá theo từng tensor

Trong quá trình lượng tử hoá theo từng tensor, một số tỷ lệ và điểm không duy nhất được dùng cho tất cả các giá trị trong tensor. Loại được lượng tử hoá trên mỗi tensor được biểu thị trong StableHLO dưới dạng:

quant.uniform scale:zero_point>

Ví dụ: !quant.uniform<i8:f32, 0.01:50>

Đây là số nguyên 8 bit (i8) dùng để lưu trữ số có dấu phẩy động 32 bit (f32) bằng cách sử dụng tỷ lệ 0.01 và điểm 0 là 50.

Lượng tử hoá theo trục

Lượng tử hoá theo trục cung cấp một phương pháp chi tiết hơn so với lượng tử hoá theo tensor. Thay vì sử dụng một tỷ lệ và điểm 0 duy nhất cho toàn bộ tensor, quá trình định lượng theo trục sẽ chỉ định các tỷ lệ và điểm 0 riêng biệt cho các lát dọc theo một phương diện cụ thể quantized_dimension của tensor. Điều này đặc biệt hữu ích khi các giá trị thay đổi đáng kể trên nhiều phương diện, giúp bảo toàn thông tin và độ chính xác tốt hơn.

Hãy xem xét một tensor t có kích thước phương diện [4, 3, 2]. Chúng ta chọn lượng tử hoá tenxơ này theo chiều thứ hai (quantized_dimension = 1). Điều này có nghĩa là chúng ta sẽ có 3 lát (vì chiều thứ hai có kích thước là 3), mỗi lát có tỷ lệ và điểm 0 riêng:

t[:, 0, :]: This slice gets scale[0] and zero_point[0].
t[:, 1, :]: This slice gets scale[1] and zero_point[1].
t[:, 2, :]: This slice gets scale[2] and zero_point[2].

Trong StableHLO, loại được lượng tử hoá theo trục được biểu thị là:

quant.uniform {scale0:zero_point0, scale1:zero_point1, ...}>

trong đó độ dài của scale:zero_point khớp với số lượng lát cắt dọc theo quantized_dimension của tensor chứa.

Ví dụ: tensor<4x3x2x!quant.uniform<i8:f32:1, {0.2:20, 0.1:10, 0.3:30}>>

Các lượt truyền lượng tử hoá trong StableHLO

StableHLO cung cấp một số lượt truyền trình biên dịch cho phép thực hiện nhiều hoạt động chuyển đổi và tối ưu hoá liên quan đến việc lượng tử hoá, giúp bạn linh hoạt trong cách xử lý các mô hình được lượng tử hoá. Các thẻ và vé này là:

stablehlo-legalize-qdq-to-quantized-op

Lượt truyền này hợp nhất một mẫu hình phổ biến trong các mô hình được định lượng, một phép toán khử lượng tử hoá theo sau là một phép toán dấu phẩy động và cuối cùng là một phép toán lượng tử hoá, thành một phép toán được định lượng duy nhất. chi tiết

stablehlo-legalize-quantized-op-to-qdq

Thẻ này có chức năng ngược lại với thẻ trước đó. Thao tác này phân tách một thao tác được lượng tử hoá thành chuỗi tương đương gồm các thao tác khử lượng tử hoá, thao tác dấu phẩy động và thao tác lượng tử hoá. chi tiết

stablehlo-legalize-quant-to-math

Lượt truyền này chuyển đổi các thao tác StableHLO trên các loại được lượng tử hoá thành các thao tác tương đương trên các loại số nguyên. Về cơ bản, hàm này triển khai phép toán số học lượng tử hoá bằng các phép toán toán học tiêu chuẩn. Việc phân tách này rất hữu ích cho những hệ thống không hỗ trợ việc lượng tử hoá một cách tự nhiên, nhưng vẫn có thể sử dụng phép toán lượng tử hoá để thể hiện ngữ nghĩa của các mô hình được lượng tử hoá. chi tiết

stablehlo-quant-legalize-to-tosa-rescale

StableHLO cung cấp khả năng hợp thức hoá các hoạt động được lượng tử hoá thành các biểu diễn tương ứng trong phương ngữ TOSA. Việc hợp pháp hoá này giúp tăng khả năng tương thích và khả năng tương tác giữa StableHLO và TOSA. Lượt truyền này chuyển đổi một cách chiến lược các thao tác được lượng tử hoá StableHLO thành một tổ hợp các thao tác StableHLO và TOSA, trong đó phương ngữ TOSA chủ yếu được dùng cho thao tác rescale. Hoạt động tosa.rescale đóng vai trò quan trọng trong việc điều chỉnh tỷ lệ và điểm 0 của các giá trị được lượng tử hoá, cho phép biểu diễn chính xác dữ liệu được lượng tử hoá trong khuôn khổ TOSA. chi tiết

tosa-rescale-legalize-to-stablehlo

Lượt truyền này ghi lại các thao tác chia tỷ lệ TOSA thành các thao tác toán học nguyên thuỷ StableHLO. Một trong những trường hợp sử dụng chính của đường truyền này là cho phép trình thông dịch StableHLO đánh giá các chương trình có chứa các thao tác điều chỉnh tỷ lệ TOSA. chi tiết

Đánh giá các chương trình được lượng tử hoá

Trình thông dịch tham chiếu StableHLO có thể thực thi hiệu quả các chương trình chứa các thao tác được lượng tử hoá. Để đạt được điều này, trước tiên, chương trình sẽ giảm xuống mức biểu thị tương đương chỉ bằng các thao tác số nguyên. Quá trình hạ cấp này bao gồm một loạt các lần truyền trình biên dịch giúp chuyển đổi chương trình trước khi diễn giải.

Về cơ bản, trình thông dịch tận dụng lượt truyền stablehlo-legalize-quant-to-math để chuyển đổi các hoạt động được lượng tử hoá thành các hoạt động tương ứng trong quá trình triển khai số học số nguyên. Lượt truyền này giới thiệu các thao tác truyền tin CHLO để xử lý phép nhân/chia theo tỷ lệ và phép cộng điểm 0. Để đảm bảo khả năng tương thích với trình thông dịch StableHLO, các thao tác CHLO này sau đó sẽ được hợp thức hoá thành các thao tác StableHLO. Thao tác này giới thiệu các thao tác liên quan đến hình dạng, sau đó được chuẩn hoá và tối ưu hoá bằng một loạt các lượt chuẩn hoá.

Sau đây là toàn bộ chuỗi các lượt truyền liên quan đến quy trình giảm này:

stablehlo-legalize-quant-to-math
chlo-legalize-to-stablehlo
canonicalize
shape-legalize-to-stablehlo
stablehlo-canonicalize-dynamism

Trường hợp kiểm thử được lượng tử hoá

StableHLO cung cấp một bộ toàn diện gồm các trường hợp kiểm thử được lượng tử hoá để xác thực tính chính xác và hành vi của các hoạt động được lượng tử hoá. Các trường hợp kiểm thử này đóng vai trò là kiểm thử đơn vị, bao gồm nhiều thao tác StableHLO trong các tình huống được định lượng.

Ví dụ điển hình về một trường hợp kiểm thử được lượng tử hoá có dạng

func.func @main() -> tensor<11xf32> {
    %operand_0 = stablehlo.constant dense<...> : tensor<11xf32>
    %operand_1 = stablehlo.constant dense<...> : tensor<11xf32>
    %golden = stablehlo.constant dense<...> : tensor<11xf32>

    %0 = stablehlo.uniform_quantize %operand_0 : (tensor<11xf32>) -> tensor<11x!quant.uniform<i8:f32, 0.3>>
    %1 = stablehlo.uniform_quantize %operand_1 : (tensor<11xf32>) -> tensor<11x!quant.uniform<i8:f32, 0.3>>

    %2 = stablehlo.add %1, %0 : tensor<11x!quant.uniform<i8:f32, 0.3>>

    %result = stablehlo.uniform_dequantize %2 : (tensor<11x!quant.uniform<i8:f32, 0.3>>) -> tensor<11xf32>

    %4 = stablehlo.custom_call @check.eq(%golden, %result) : (tensor<11xf32>, tensor<11xf32>) -> tensor<i1>

    return %3 : tensor<11xf32>
  }

và bao gồm:

Những trường hợp kiểm thử này có giá trị đối với:

  • Xác thực việc lượng tử hoá StableHLO: Đảm bảo rằng hành vi lượng tử hoá của các hoạt động StableHLO phù hợp với kết quả dự kiến.
  • Xác thực chéo: So sánh hành vi của lượng tử hoá StableHLO với các cách triển khai hoặc khung khác.
  • Gỡ lỗi và phát triển: Hỗ trợ phát triển và gỡ lỗi các tính năng lượng tử hoá hoặc tối ưu hoá mới.