Optimum项目中的Sentence Transformers模型导出ONNX时torch_dtype参数失效问题分析

问题背景

在Hugging Face的Optimum项目中，用户发现当尝试将Sentence Transformers模型导出为ONNX格式时，指定的数据类型参数（如bfloat16）会被忽略，导致导出的模型始终使用默认的float32精度。这个问题主要影响那些希望优化模型大小和推理性能的用户，特别是在资源受限的环境中。

技术细节

问题出现在Optimum的模型导出任务处理代码中。当用户通过命令行工具指定--dtype bf16参数时，预期的行为是模型应该以bfloat16精度导出到ONNX格式。然而，当前实现中这个参数没有被正确传递到SentenceTransformer模型的构造函数中。

具体来说，在optimum/exporters/tasks.py文件的模型初始化部分，虽然代码接收了torch_dtype参数，但没有将其传递给SentenceTransformer的构造函数。正确的做法应该是通过model_kwargs字典将这个参数传递给模型。

影响范围

这个问题会影响所有使用Optimum工具导出Sentence Transformers模型到ONNX格式的用户，特别是：

1. 需要降低模型精度以减少内存占用的场景
2. 希望在特定硬件上获得更好性能的用户
3. 需要控制模型导出精度的研究者和工程师

解决方案

修复方案相对简单直接，只需在模型初始化时正确传递torch_dtype参数。具体修改是在SentenceTransformer模型初始化时，将数据类型参数通过model_kwargs字典传递：

model = model_class(
    model_name_or_path,
    device=device,
    cache_folder=cache_folder,
    token=token,
    trust_remote_code=trust_remote_code,
    model_kwargs={"torch_dtype": torch_dtype}
)

后续建议

对于遇到此问题的用户，可以：

1. 等待官方合并修复补丁
2. 临时手动修改本地Optimum安装中的相关代码
3. 在导出后使用ONNX工具进行精度转换

这个问题的修复将增强Optimum工具链的完整性和可靠性，确保用户能够精确控制模型导出的各项参数，满足不同场景下的精度需求。