XLA 设备的量化运算（实验性功能）

---

本文档概述了如何利用量化运算在 XLA 设备上启用量化。

XLA 量化运算为量化运算（例如，块级 int4 量化矩阵乘法）提供了高级抽象。这些作类似于 CUDA 生态系统中的量化 CUDA 内核（示例），在 XLA 框架中提供类似的功能和性能优势。

注意：目前，这被归类为实验性功能。这是 API 的具体内容 将在下一个 （2.5） 版本中更改。

如何使用：

XLA 量化运算可用作 ，或将 .这两个选项使模型开发人员能够灵活地选择将 XLA 量化运算集成到其解决方案中的最佳方式。torch optorch.nn.Moduletorch.op

两者 都与 兼容。torch opnn.Moduletorch.compile( backend='openxla')

在模型代码中调用 XLA 量化运算

用户可以像调用其他常规 PyTorch 运算一样调用 XLA 量化运算。这为将 XLA 量化运算集成到其应用程序中提供了最大的灵活性。量化运算在 Eager 模式和 Dynamo 下均可工作，具有常规 PyTorch CPU 张量和 XLA 张量。

注意请检查量化运算的文档字符串，了解量化权重的布局。

import torch
import torch_xla.core.xla_model as xm
import torch_xla.experimental.xla_quantized_matmul

N_INPUT_FEATURES=10
N_OUTPUT_FEATURES=20
x = torch.randn((3, N_INPUT_FEATURES), dtype=torch.bfloat16)
w_int = torch.randint(-128, 127, (N_OUTPUT_FEATURES, N_INPUT_FEATURES), dtype=torch.int8)
scaler = torch.randn((N_OUTPUT_FEATURES,), dtype=torch.bfloat16)

# Call with torch CPU tensor (For debugging purpose)
matmul_output = torch.ops.xla.quantized_matmul(x, w_int, scaler)

device = xm.xla_device()
x_xla = x.to(device)
w_int_xla = w_int.to(device)
scaler_xla = scaler.to(device)

# Call with XLA Tensor to run on XLA device
matmul_output_xla = torch.ops.xla.quantized_matmul(x_xla, w_int_xla, scaler_xla)

# Use with torch.compile(backend='openxla')
def f(x, w, s):
  return torch.ops.xla.quantized_matmul(x, w, s)

f_dynamo = torch.compile(f, backend="openxla")
dynamo_out_xla = f_dynamo(x_xla, w_int_xla, scaler_xla)

通常将量化运算包装到模型开发人员模型代码中的自定义中：nn.Module

class MyQLinearForXLABackend(torch.nn.Module):
  def __init__(self):
    self.weight = ...
    self.scaler = ...

  def load_weight(self, w, scaler):
    # Load quantized Linear weights
    # Customized way to preprocess the weights
    ...
    self.weight = processed_w
    self.scaler = processed_scaler


  def forward(self, x):
    # Do some random stuff with x
    ...
    matmul_output = torch.ops.xla.quantized_matmul(x, self.weight, self.scaler)
    # Do some random stuff with matmul_output
    ...

模块交换

或者，用户也可以使用包装 XLA 量化运算并在模型代码中进行模块交换的 the：nn.Module

orig_model = MyModel()
# Quantize the model and get quantized weights
q_weights = quantize(orig_model)
# Process the quantized weight to the format that XLA quantized op expects.
q_weights_for_xla = process_for_xla(q_weights)

# Do module swap
q_linear = XlaQuantizedLinear(self.linear.in_features,
                              self.linear.out_features)
q_linear.load_quantized_weight(q_weights_for_xla)
orig_model.linear = q_linear

支持的量化运算：

矩阵乘法

权重量化类型	Activation Quantization Type（激活量化类型）	D型	支持
每通道 （SYM/ASYM）	不适用	W8A16 系列	是的
每通道 （SYM/ASYM）	不适用	W4A16 系列	是的
每通道	每个令牌	W8A8 系列	不
每通道	每个令牌	W4A8 系列	不
按块 （SYM/ASYM）	不适用	W8A16 系列	是的
按块 （SYM/ASYM）	不适用	W4A16 系列	是的
块	每个令牌	W8A8 系列	不
块	每个令牌	W4A8 系列	不

注意 Weight in -bit， Activation in -bit。如果为 4 或 8，则引用 。16 表示格式。W[X]A[Y]XYX/Yint4/8bfloat16

嵌入

待添加