PyTorch AO项目中的量化工作流配置迁移：从可调用对象到配置对象

概述

在PyTorch AO（torchao）项目中，量化工作流的配置方式正在经历一次重要的架构演进。从0.9.0版本开始，项目团队决定将quantize_工作流的配置方式从使用可调用对象（Callables）迁移到直接使用配置对象（config objects）。这一变更旨在更好地与生态系统其他部分对齐，支持配置实例化后的检查能力，并消除常见的混淆来源。

变更内容详解

配置方式的演进

在0.8.0及之前版本中，quantize_函数的第二个参数是一个可调用对象：

def quantize(
    model: torch.nn.Module,
    apply_tensor_subclass: Callable[[torch.nn.Module], torch.nn.Module],
    ...,
): ...

从0.9.0版本开始，该参数将支持两种形式：

def quantize(
    model: torch.nn.Module,
    config: Union[AOBaseConfig, Callable[[torch.nn.Module], torch.nn.Module]],
    ...,
): ...

而在未来的0.10.0或更高版本中，将完全过渡到配置对象：

def quantize(
    model: torch.nn.Module,
    config: AOBaseConfig,
    ...,
): ...

工作流配置名称变更

所有工作流配置的名称从蛇形命名法（snake_case）改为驼峰命名法（CamelCase），同时保留了旧名称作为别名以确保向后兼容：

旧名称（保持可用）	新名称（推荐使用）
int4_weight_only	Int4WeightOnlyConfig
float8_dynamic_activation_float8_weight	Float8DynamicQuantizationFloat8WeightConfig
...	...

迁移影响分析

不受影响的用户

大多数现有用户不会受到影响，特别是那些使用位置参数传递配置的用户：

quantize_(model, int8_weight_only(group_size=128))

这种语法将继续有效，用户可以选择在合适的时间迁移到新的配置名称：

quantize_(model, Int8WeightOnlyConfig(group_size=128))

需要调整的用户

少数使用关键字参数传递配置的用户需要修改调用方式：

# 旧方式（需要修改）
quantize_(model, tensor_subclass_inserter=int8_weight_only(group_size=128))

# 新方式
quantize_(model, config=int8_weight_only(group_size=128))

开发者注意事项

为quantize_API开发新工作流的开发者需要使用新的配置系统。新系统基于抽象基类AOBaseConfig和具体的数据类配置：

@dataclass
class Int4WeightOnlyConfig(AOBaseConfig):
    group_size: int = 128
    ...

技术实现细节

旧实现方式

旧系统通过返回可调用对象的方式工作：

def int4_weight_only(group_size: int, ...) -> Callable:
    def new_callable(weight: torch.Tensor):
        # 配置通过局部变量捕获
        ...
    return _get_linear_subclass_inserter(new_callable)

新实现架构

新系统采用更结构化的方式：

1. 定义配置基类
2. 创建具体的数据类配置
3. 注册模块转换处理器

@register_quantize_module_handler(Int4WeightOnlyConfig)
def _int4_weight_only_transform(
    module: torch.nn.Module, 
    config: Int4WeightOnlyConfig,
) -> torch.nn.Module:
    # 实现具体的转换逻辑
    ...

迁移路线图

项目团队制定了清晰的迁移计划：

1. 0.9.0版本：启用新语法，同时开始弃用旧语法
2. 0.10.0或更高版本：完全移除旧语法支持

在整个过程中，旧名称将作为新名称的别名永久保留，确保现有代码库的兼容性。

对生态系统的影响

项目团队已经验证了主要合作伙伴（如Hugging Face Transformers、SGLANG和Diffusers）的集成仍然正常工作，确保生态系统的平稳过渡。

这一架构改进将为PyTorch量化生态系统带来更清晰、更可维护的代码结构，同时为未来的功能扩展奠定更坚实的基础。