机器学习软件与 pickle 是怎么回事?
原文由 Nelson Elhage 于 发布,订阅该博客
我做了很多年软件工程师,完全是机器学习领域的门外汉,只是出于好奇,偶尔在外围与它有些接触。在那段时间里,每当看到 Python 机器学习生态中普遍使用 Python 的 pickle,我一贯的感受就是震惊——说实话,还有不屑。
除了严重的安全问题1,使用 pickle 做序列化还非常脆弱,随着代码演进、库和 Python 版本的升级,会带来各种噩梦。在我做软件工程师的职业生涯里,我一度近乎认定——可以说相当准确地——pickle 模块没有任何合理的使用场景,它纯粹是一个诱人的陷阱;是那个更年轻、更冲动、也更天真的生态所犯下的一个错误。
不过,过去几年里我开始在机器学习领域全职工作,这让我对这个生态有了新的视角。我切身体会到了机器学习软件为何会呈现出现在的样子,也理解了 pickle 到底在为用户解决哪些问题。我依然不喜欢它,但我对这个领域的问题有了更多的理解和共情,也不再认为这个问题是“小菜一碟”,或仅仅是无知、懒惰所致。这篇文章就是想分享一下我的这些体会,试着在这两个世界之间架起一点桥梁。
做研究,而不是做软件
归根结底,我最想传达的核心观点是:机器学习生态中的绝大多数软件,首先是为了服务研究人员而设计的,而不是为了服务软件工程师。进一步说,(至少在当下的语境里)研究本身具有——甚至可以说就定义于——以下几个相互关联的特点:
- 研究的首要产出是知识,而不是软件产物。研究团队写软件是为了回答研究问题、加深自己、团队乃至整个领域对某个方向的理解,而不是为了交付软件工具或解决方案。
- 研究人员往往事先并不知道哪些想法会成功、哪些会有用,因此他们优先考虑的是高效地尝试大量想法,并且默认其中大多数都会失败。
作为上述特点的推论,研究人员(尤其是在机器学习这样的领域!)可能会写下大量代码,但这些代码中的绝大多数很快就会被丢弃,这也完全正常。失败的实验通常会被直接抛弃,即使是成功的实验,也往往在论文写完后就不再使用了。即便后续工作要沿用某个成功的思路,也常见的做法是为新项目重写代码,因为此时对问题的理解更深了,而且要面对的是新的问题或研究课题。
相比典型的软件工程场景,这种环境下的权衡与便利性要远更有利于 pickle(相对于其他序列化方案):
- 大多数代码自始至终都只由同一小撮人(研究者及其团队)在单一环境中编写和运行;往往几乎不区分“开发环境”和“生产环境”,也没有外部用户、外部数据源或外部交互。这些因素叠加在一起,使得
pickle序列化在实践中的安全风险要低得多(但绝非为零!) - 能够直接拿起一个数据结构——甚至一段代码——就地序列化,而不必停下来考虑模式、格式或数据演进,有着巨大的好处;我们看重的是多快能拿到结果,而花精力去琢磨模式、序列化格式,或在磁盘格式与内存格式之间做转换,完全是额外负担。
- 由于大多数代码很快就会被丢弃,
pickle的脆弱性以及处理代码变更和版本升级的难题,就显得没那么突出了。
一个研究小记
为了让上述权衡更具体一些,我想用一种程式化的方式,描述一位从事机器学习研究的研究人员,试着展示她会在哪些具体场景下用到 pickle,以及在我看来其中的利弊——这些看法来自我自己做机器学习研究、以及支持(远比我资深得多的!)机器学习研究人员的经验。
假设一位研究人员正在尝试一种新的深度学习模型架构,或对现有架构做变体。她的架构会有大量配置选项和超参数:层数、每层的类型、各种向量的维度、在何处以及如何对激活值做归一化、使用哪种非线性函数,等等。模型的许多组件会是机器学习框架提供的标准层,但研究人员也会在其中插入一些新颖的逻辑片段。
我们的研究人员需要一种方式来描述一个具体的模型——即这些设置的特定组合——并能将其序列化、之后再重新加载。她出于几个相关的原因需要做到这一点:
- 她很可能拥有一套包含 GPU 或其他加速器的计算集群来跑任务。她需要一种方式,把模型描述提交给集群上运行的代码,让集群去跑她的模型。
- 在模型训练期间,她需要以可重新加载、续训的方式保存训练进度的快照,以防硬件故障或任务被抢占。
- 模型训练完成后,研究人员会想要再次加载它们(可能既包括最终的快照,也包括部分训练过程中的检查点),以便进行评估和实验。
我们可以设想许多不同的解决方案!我们可以写一个 protobuf 或 JSON 的模型描述,再写代码把它转换成真正的模型实例。我们甚至可以设想把模型表达为“一组命令行参数”,再解析成模型实例。
然而,这里有一个能一举解决所有这些问题、而且前期几乎不需要额外工作量的方案:
- “一个具体模型”的定义就是“一个实现了该模型的对象”(很可能是一个
torch.nn.Module的子类) - 研究人员通过一个临时性的研究脚本,以某种即兴的方式构造这些对象,并在开展新实验的过程中不断编辑、演进这个脚本
- 然后她直接用
pickle把整个对象序列化,并将序列化结果作为存储版本,用于后续分析或在远程机器上加载执行。
这种做法不需要任何繁琐操作,对研究代码或实现如何组织也不要求特定的规范或模式,相比“仅仅是构造并运行一个模型”,几乎没有额外的心智负担。它对标准组件和自定义组件一视同仁:研究人员可以在类层次结构深处的任意位置做小幅改动,也可以替换整个模型架构,而这两者都不需要过多考虑或修改序列化系统。
到这一步,我们的研究人员很可能会写一些脚本来生成各种超参数的笛卡尔积组合,把它们全部调度到集群上,然后就去吃午饭(或者度周末去了),让它们慢慢训练。
这些 pickle 文件无疑是非常脆弱的!它们把模型的完整实现细节都序列化了进去,随着研究人员添加新的设置和参数、重构实验,很容易就会破坏旧的实验和模型。但话说回来,大多数实验本就不会得出有趣的结果,之后也根本不需要再看。对于少数有价值的实验,只要记下每次实验对应的 git 提交,需要时回退代码检出即可。而且如果研究人员需要更持久地保留某几个特定模型,通常也可以在某个 __setstate__ 方法里做几处针对性的修补来恢复。
话虽如此,如果这批实验总体上是成功的,我们需要沿着这条思路继续深挖,尝试更多变体,并且为了对比而更长期地保留模型,那么这种脆弱性就开始让人头疼了,于是我们的研究人员可能会考虑重构。一个直接的办法是把模型的“描述”与实现分离开来:把模型的配置、设置和超参数归拢到一个 class ModelConfig: 中;真正被序列化的(依然通过 pickle)是这个类,而研究人员则手写维护一个 build(cfg: ModelConfig) -> Model 函数,负责根据配置实例化出完整的模型。
这样一来,研究人员只需要关心这个类的向前和向后兼容性,而实现代码则可以相对自由地重构,只要同步重构从 Model 对象构造 ModelConfig 的那部分代码即可。
但值得注意的是,我们的研究人员(和/或她的团队)仍然在做实验,不断尝试新东西,迭代不止,此时“多快能做实验”,甚至更重要的是“做实验的心智成本”压倒一切。因此,虽然这个类上的许多字段会是规规矩矩的整数、布尔值乃至枚举,但要为每一个新实验都去添加支持,还是有点繁琐。
于是,或许我们的研究人员会留下一些“后门”。举一个最简单的例子:在这个领域,尝试在全连接层之后替换非线性函数或“激活函数”仍然相当流行(我自己也没少干)。
我们该如何在配置中存储一个激活函数(它很可能就是一个类型为 Tensor -> Tensor 的 Python 函数)呢?
我们可以在某处注册一个 dict[str, Callable],按名称来存储函数……但话说回来,Python 的模块命名空间本身不就是这样一个字典吗?所以如果我们用的是 pickles,大概就可以直接把函数存到配置里了。
受此启发,我们的研究人员更进一步;她甚至可以在配置里存入一个完整的 build_xxxx: Callable[[Config], Module | None] 重写,让配置以任意方式覆盖任意部分的配置!
这类特性让大部分原有的脆弱性问题又回来了;但重要的是,它们让研究团队得以在两者之间取得平衡:一边是那些更稳定、更核心的配置,以“正规”参数的形式成为一等公民;另一边则是实验中泡沫般快速更迭的部分,那里(几乎)没有什么能长久存在,迭代速度就是一切。
范围蔓延
在我个人看来,只要项目确实只在可信环境中用于内部研究,以上我描述的一切都“基本可以接受”,虽然有点让人不舒服、也有风险,前提是项目确实只在可信环境中供内部研究使用。
但更大的问题在于,一旦你开始以这种方式常规地使用 pickle,就会习惯性地用它来解决所有的序列化问题,哪怕在那些 pickle 远为棘手的场景中也是如此。不仅研究人员会这样做,这种模式还会渗透到工具和整个生态中。
例如,PyTorch 就把 pickle 文件作为其默认序列化格式的一部分(通过 torch.load 和 torch.save),尽管在大多数情况下,它们真正需要的只是一些“朴素”的数据结构——字符串和 Tensor 构成的列表和字典。PyTorch 甚至还有一个只支持少数内置类型的定制版 C++ pickle 加载器,可以通过加载器的 weights_only 参数启用;但为了向后兼容,这个标志默认是 False,而且(在没有攻击者的情况下)启用它只会带来破坏;因此它很少被使用,每个 PyTorch 权重文件都是在向用户敞开大门,邀请他们使用这个危险的接口,让自己暴露于风险之中。
该怎么办?
我写这篇文章的主要目的,是想呈现一下“内部视角”——为什么我在机器学习生态中看到如此多的 pickle 使用,以及它在为用户解决哪些问题。话虽如此,随着机器学习变得越来越重要,用户和部署场景也越来越多样,我预见不安全的 pickle 使用所引发的问题会变得愈发重要和紧迫,因此我想就未来的出路提供一些思考。
更好的数据格式与接口
首先,我认为我们(在这里我主要指的是那些拥有软件工程背景和技能、和/或在机器学习领域中主要承担软件工程角色的人)应该继续投入,打造 pickle 的高质量替代方案,从而为研究人员提供更多、更好的选择。如果做得好,我们甚至可以在许多场景下为研究人员提供比 pickles 更好的体验,同时具备更好的安全性。
在此我不得不指出,围绕 XLA 构建的 Google 机器学习生态——Tensorflow 和 jax 及其各种库和框架——在很大程度上避免了这里所描述的问题;它们倾向于使用基于 protocol buffer 和 numpy 格式的序列化格式,pickle 的泛滥程度要低得多。一方面,我把它们视为这种生态是可行的某种存在性证明;另一方面,我倾向于认为 PyTorch 之所以能占据主导地位,正是因为它总体上为研究人员提供了更优的使用体验,尤其是在上手和易于修改、探索方面,因此我不太愿意把 Google 的那套库当作理想的解决方案来推荐。
Huggingface 的 safetensors 库是这个领域进展的一个绝佳例子,它为安全、高效地序列化张量数据提供了坚实的基础。在我看来,它最好被看作一个略为底层的原语;它支持对形如 dict[str, Tensor] 的对象进行序列化和恢复;但研究人员往往想要的是嵌套字典,或希望在同一个包中同时存储权重和元数据;我很期待看到在 safetensors 之上,涌现出更符合人体工学、更对研究人员友好的格式和 API。
研究人员往往也想序列化 Python 类树以及张量数据;也许 Python 类型注解近年来的出现和普及,以及基于其之上的 attrs 和 dataclasses 等 API,为一个由 Python 类型注解和少量额外元数据驱动的库提供了可能(或许可以基于 cattrs 构建),从而能够或多或少直接地序列化和反序列化[某些]带类型注解的 Python 类树。
限制 pickle
然而,无论好坏,我们今天已经有了大量 pickle 文件和在使用 pickle 文件的库,而且这个问题还在不断扩大。在我们努力提供替代方案、并鼓励项目迁移到这些替代方案的同时,这里有一个关于减少危害或缓解风险的推测性想法:
首先,我在想 Python 是否可以支持一个全局的“nopickle”模式来彻底禁用 pickle 模块。这样的标志至少能让部署 Python 的运维人员轻松判断,他们在运行时加载不受信任的 pickle 时是否面临风险。
然而,如果那样过于严格,那么为 pickle 的使用添加一些限制选项又如何呢?
Pickle 清单
特别是,我对这样一个想法很感兴趣:要求所有对 pickle 的使用都声明一个可用于解码的合法二进制对象列表,很可能以 sha256 或其他密码学哈希来命名。我们可以设想,既支持在 pickle.load 调用处(或某个新 API 处)指定预期对象,也支持在某种清单文件中带外指定,甚至可以指定合法的(调用栈,blob ID)元组。任何未指定此数据,或与允许的哈希不匹配的 pickle 加载都会被拒绝。
这个方案之所以让我感兴趣,是因为它或许能让部署机器学习软件的人——只需一次——收集所有正在被加载的模型权重文件或配置文件的清单,然后“冻结”这个集合,确保在运行时不会遇到任何新的 pickle。
这反过来又在很大程度上把管理 pickle 使用的问题,归约为管理第三方依赖的问题;可以肯定,这绝不是一个已解决的问题,但它本来就是你无论如何都要面对的问题。
具体来说,要部署一个新的第三方机器学习模型,你通常需要同时安装模型实现(以 Python 模块或脚本的形式)和模型权重。这两者如果带有恶意,都会带来安全风险;但 pickle 带来的风险在某种意义上更糟,因为即使你审查并锁定了模型代码的某个特定版本,pickle 文件往往仍会在运行时被下载,这意味着你依然暴露在来自系统外部的远程代码执行风险之下。
如果我们能保证那段代码只会加载特定的 pickle 文件——而我们在首次决定信任该依赖时就可以一次性审计这些文件——那么我们就让 pickle 的用例“看起来像”安装第三方代码的问题,即一旦我们获取并锁定某个版本,至少不会再有额外的运行时攻击面。
我们可以设想把这些想法进一步复杂化,引入一套运行时配置系统;极端情况下,我们可能会设想用这样一些相当复杂的 API 来控制 pickle 的使用:
with pickle.no_pickles() as token:
pickle.loads(...) # NoPickleException
with pickle.allow_pickles(token):
pickle.loads(...)
with pickle.allow_pickles(token, only_objects=[
"sha256:some sha256",
]):
pickle.loads(buf) # hashes buf, only proceeds if it matches
# You must pass the same token returned by `no_pickles` to
# `allow_pickle`; this prevents libraries from trivially blanket
# calls to renable `pickle`
with pickle.allow_pickles(bad_token): # BadTokenException
...当然,这类机制永远不会是完美的;库总是可以通过各种手段绕过它们。极端情况下,库总是可以自带一个纯 Python 实现的 pickle,而不受任何这些限制。然而,这只是库可能(有意或无意地)植入后门的无数种方式之一;这个提议的目标并不是要解决供应链安全问题,而是要缓解 pickle 带来的特定风险,并且如上所述,把风险从运行时转移到包安装时。我们会希望,那些公然绕过标准安全机制的库会被社区注意到,并被标记为不可信2。
这些功能可以在解释器之外,以一个对 pickle 模块进行 monkey patch 的库的形式先行原型化,但从标准化和推广的角度来看,最终将其内置到解释器中或许是有价值的。
随机一篇博客
评论
登录后参与讨论