LoRA-Scripts项目中的PyTorch 2.6安全策略与VAE模型加载问题解析

问题背景

在LoRA-Scripts项目中使用PyTorch 2.6版本进行模型训练时，用户遇到了无法加载自定义VAE模型的问题。错误信息显示这是由于PyTorch 2.6版本引入的安全策略变更导致的，具体表现为默认启用了weights_only=True参数，限制了模型加载时的全局变量范围。

技术分析

PyTorch 2.6版本为了提高安全性，对torch.load()函数的行为做出了重要修改：

1. 默认参数变更：weights_only参数从默认值False改为True，这意味着PyTorch默认只允许加载包含张量数据的模型文件，禁止执行任意代码。

2. 安全限制：当遇到包含不被允许的全局变量（如示例中的pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint）时，会抛出UnpicklingError异常。

3. 错误提示：PyTorch提供了两种解决方案：

   • 显式设置weights_only=False（但需确保模型来源可信）
   • 使用torch.serialization.add_safe_globals()或上下文管理器将特定类加入白名单

解决方案比较

针对这一问题，社区和项目开发者提出了几种不同的解决方案：

1. 环境变量法：设置TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD = "1"可以强制PyTorch使用旧版加载行为。这种方法简单直接，但可能带来潜在安全风险。

2. 模型格式转换：将模型转换为safetensors格式。这是最安全的解决方案，因为safetensors格式专门设计为只包含张量数据，不执行任何代码。

3. 代码修改：如ComfyUI项目所做的，通过捕获异常并尝试不同加载策略的方式实现兼容性。这种方法较为复杂，但能提供更好的用户体验。

最佳实践建议

对于LoRA-Scripts项目用户，我们推荐以下处理流程：

1. 优先转换模型格式：使用工具将VAE模型转换为safetensors格式，这是最安全可靠的长期解决方案。

2. 临时解决方案：在确保模型来源可信的情况下，可以使用环境变量临时解决问题，但应注意这仅适用于开发和测试环境。

3. 版本兼容性：考虑在项目文档中明确标注PyTorch版本要求，或提供自动检测和处理的代码逻辑。

技术影响评估

这一变更对深度学习社区的影响较大：

1. 安全性提升：有效防止了通过恶意模型文件执行任意代码的风险。

2. 兼容性挑战：许多现有模型文件需要更新或转换格式才能正常工作。

3. 开发习惯改变：开发者需要更加注意模型来源的可信度，并考虑采用更安全的模型存储格式。

结论

PyTorch 2.6的安全策略变更代表了深度学习框架向更安全方向发展的趋势。虽然短期内可能带来一些兼容性问题，但从长远来看，推动使用safetensors等更安全的模型格式对社区是有益的。LoRA-Scripts项目的用户应逐步适应这一变化，优先考虑将模型转换为安全格式，而不是依赖临时解决方案。