XLA设备的量化操作（实验性功能）¶
本文档概述了如何利用量化操作在XLA设备上启用量化。
XLA量化操作提供了一种高级抽象，用于量化操作（例如，块状int4量化矩阵乘法）。这些操作类似于CUDA生态系统中的量化CUDA内核（example），在XLA框架内提供了类似的功能和性能优势。
注意： 目前这被归类为实验性功能。其API细节将在下一个（2.5）版本中发生变化。
如何使用:¶
XLA量化操作可以用作torch op，或者作为包装torch.op的torch.nn.Module。这两种选择为模型开发者提供了灵活性，以选择最佳方式将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
