XLA 量化操作（实验性功能）¶
本文档概述了如何利用量化操作在 XLA 设备上启用量化。
XLA 量化操作为量化操作（例如，分块 int4 量化矩阵乘法）提供了一个高级抽象。 这些操作类似于 CUDA 生态系统中的量化 CUDA 内核（示例），在 XLA 框架内提供类似的功能和性能优势。
注意： 目前这被归类为实验性功能。其 API 规范将在下一个 (2.5) 版本中更改。
如何使用：¶
XLA 量化操作可以用作 torch op，或包装 torch.op 的 torch.nn.Module。 这 2 个选项使模型开发者能够灵活地选择将 XLA 量化操作集成到其解决方案中的最佳方式。
torch op 和 nn.Module 都与 torch.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 量化操作的 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
