(原型) FX 图模式量化用户指南¶
作者: Jerry Zhang
FX 图模式量化需要符号可跟踪模型。我们使用 FX 框架将符号可跟踪的 nn.Module 实例转换为 IR,并在 IR 上执行量化传递。请在 PyTorch 讨论论坛 上发布有关符号跟踪模型的问题
量化仅对模型的符号可跟踪部分有效。使用符号跟踪值的 -if 语句/for 循环等等 - 数据依赖控制流是常见的模式,但不受支持。如果您的模型不是端到端符号可跟踪的,您可以使用以下几种选项,仅对模型的一部分启用 FX 图模式量化。您可以组合使用这些选项
- 不可跟踪代码不需要量化
仅符号跟踪需要量化的代码
跳过符号跟踪不可跟踪的代码
- 不可跟踪代码需要量化
重构您的代码使其符号可跟踪
编写您自己的观察和量化子模块
如果不需要量化不可跟踪代码,我们可以使用以下两种选项来运行 FX 图模式量化
仅符号跟踪需要量化的代码¶
当整个模型无法进行符号追踪,但我们要量化的子模块可以进行符号追踪时,我们可以仅对该子模块运行量化。
之前
class M(nn.Module):
def forward(self, x):
x = non_traceable_code_1(x)
x = traceable_code(x)
x = non_traceable_code_2(x)
return x
之后
class FP32Traceable(nn.Module):
def forward(self, x):
x = traceable_code(x)
return x
class M(nn.Module):
def __init__(self):
self.traceable_submodule = FP32Traceable(...)
def forward(self, x):
x = self.traceable_code_1(x)
# We'll only symbolic trace/quantize this submodule
x = self.traceable_submodule(x)
x = self.traceable_code_2(x)
return x
量化代码
qconfig_mapping = QConfigMapping().set_global(qconfig)
model_fp32.traceable_submodule = \
prepare_fx(model_fp32.traceable_submodule, qconfig_mapping, example_inputs)
注意,如果需要保留原始模型,则必须在调用量化 API 之前自行复制它。
跳过对不可追踪代码的符号追踪¶
当模块中存在一些不可追踪的代码,并且这部分代码不需要进行量化时,可以将这部分代码分解成一个子模块,并跳过对该子模块的符号追踪。
之前
class M(nn.Module):
def forward(self, x):
x = self.traceable_code_1(x)
x = non_traceable_code(x)
x = self.traceable_code_2(x)
return x
之后,不可追踪部分移到一个模块并标记为叶子
class FP32NonTraceable(nn.Module):
def forward(self, x):
x = non_traceable_code(x)
return x
class M(nn.Module):
def __init__(self):
...
self.non_traceable_submodule = FP32NonTraceable(...)
def forward(self, x):
x = self.traceable_code_1(x)
# we will configure the quantization call to not trace through
# this submodule
x = self.non_traceable_submodule(x)
x = self.traceable_code_2(x)
return x
量化代码
qconfig_mapping = QConfigMapping.set_global(qconfig)
prepare_custom_config_dict = {
# option 1
"non_traceable_module_name": "non_traceable_submodule",
# option 2
"non_traceable_module_class": [MNonTraceable],
}
model_prepared = prepare_fx(
model_fp32,
qconfig_mapping,
example_inputs,
prepare_custom_config_dict=prepare_custom_config_dict,
)
如果需要对不可追踪的代码进行量化,我们有以下两种选择
重构代码使其可进行符号追踪¶
如果代码易于重构并使其可进行符号追踪,我们可以重构代码,并删除 python 中对不可追踪构造的使用。
有关符号追踪支持的更多信息,请参见 此处.
之前
def transpose_for_scores(self, x):
new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
x = x.view(*new_x_shape)
return x.permute(0, 2, 1, 3)
这无法进行符号追踪,因为在 x.view(*new_x_shape) 中不支持解包,但是,由于 x.view 也支持列表输入,因此可以轻松删除解包操作。
之后
def transpose_for_scores(self, x):
new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
x = x.view(new_x_shape)
return x.permute(0, 2, 1, 3)
这可以与其他方法结合使用,量化代码取决于模型。
编写自己的观测和量化子模块¶
如果不可追踪的代码无法重构为可进行符号追踪,例如它包含一些无法消除的循环,例如 nn.LSTM,则需要将不可追踪的代码分解成一个子模块(在 fx 图模式量化中称为 CustomModule),并定义子模块的观测版本和量化版本(在训练后静态量化或量化感知训练中用于静态量化)或定义量化版本(在训练后动态量化和仅权重量化中)。
之前
class M(nn.Module):
def forward(self, x):
x = traceable_code_1(x)
x = non_traceable_code(x)
x = traceable_code_1(x)
return x
之后
1. 将 non_traceable_code 分解为 FP32NonTraceable 不可追踪逻辑,并封装在一个模块中
class FP32NonTraceable:
...
2. 定义 FP32NonTraceable 的观测版本
class ObservedNonTraceable:
@classmethod
def from_float(cls, ...):
...
3. 定义 FP32NonTraceable 的静态量化版本,以及一个类方法“from_observed”用于将 ObservedNonTraceable 转换为 StaticQuantNonTraceable
class StaticQuantNonTraceable:
@classmethod
def from_observed(cls, ...):
...
# refactor parent class to call FP32NonTraceable
class M(nn.Module):
def __init__(self):
...
self.non_traceable_submodule = FP32NonTraceable(...)
def forward(self, x):
x = self.traceable_code_1(x)
# this part will be quantized manually
x = self.non_traceable_submodule(x)
x = self.traceable_code_1(x)
return x
量化代码
# post training static quantization or
# quantization aware training (that produces a statically quantized module)v
prepare_custom_config_dict = {
"float_to_observed_custom_module_class": {
"static": {
FP32NonTraceable: ObservedNonTraceable,
}
},
}
model_prepared = prepare_fx(
model_fp32,
qconfig_mapping,
example_inputs,
prepare_custom_config_dict=prepare_custom_config_dict)
校准 / 训练(未显示)
convert_custom_config_dict = {
"observed_to_quantized_custom_module_class": {
"static": {
ObservedNonTraceable: StaticQuantNonTraceable,
}
},
}
model_quantized = convert_fx(
model_prepared,
convert_custom_config_dict)
训练后动态/仅权重量化在这两种模式下,我们不需要观测原始模型,因此只需要定义量化模型。
class DynamicQuantNonTraceable: # or WeightOnlyQuantMNonTraceable
...
@classmethod
def from_observed(cls, ...):
...
prepare_custom_config_dict = {
"non_traceable_module_class": [
FP32NonTraceable
]
}
# The example is for post training quantization
model_fp32.eval()
model_prepared = prepare_fx(
model_fp32,
qconfig_mapping,
example_inputs,
prepare_custom_config_dict=prepare_custom_config_dict)
convert_custom_config_dict = {
"observed_to_quantized_custom_module_class": {
"dynamic": {
FP32NonTraceable: DynamicQuantNonTraceable,
}
},
}
model_quantized = convert_fx(
model_prepared,
convert_custom_config_dict)
你也可以在测试中的 test_custom_module_class
中找到自定义模块的示例,位于 torch/test/quantization/test_quantize_fx.py
。