AI安全大瓜！某互联网公司大模型训练时被员工投毒

关注兰花豆，探讨网络安全

10月18日，微信群疯传某互联网大厂AI岗位员工因对计算资源分配不满而产生报复心理，疑似对大模型训练参数投毒，利用Hugging face中pickle.load()函数在反序列化过程中可以执行任意数据（包括恶意代码）来达到攻击目的，导致其训练成果无法达到预期，损失将超过千万美元。

Hugging face介绍

首先让我们了解一下Hugging face，号称机器学习领域的“GitHub”，是一家专注于自然语言处理（NLP）和机器学习模型的技术公司。它提供了一个广泛使用的开源库，特别是用于处理和部署基于预训练模型的解决方案。Hugging Face通过其模型库和工具、数据集，简化了NLP和其他机器学习任务的开发、训练、评估和部署流程。

Transformers介绍

Hugging Face Transformers是由Hugging Face创建的深度学习开源框架（注意Transformers和Transformer区别）。它提供API和工具来下载最先进的预训练模型，并进一步调整它们以最大限度地提高性能。这些模型支持不同模式下的常见任务，例如自然语言处理、计算机视觉、音频和多模式应用程序。Transformers支持在PyTorch、TensorFlow和JAX上的互操作性.，这给在模型的每个阶段使用不同的框架带来了灵活性。

.ckpt文件介绍

.ckpt文件通常是用于保存模型检查点（checkpoint）的文件。这些文件存储了模型的权重、优化器状态、训练状态等信息，以便可以恢复模型并继续训练或用于推理。.ckpt文件的作用如下：

保存模型的训练进度：当你在训练深度学习模型时，模型的权重和训练状态会被定期保存为检查点文件，以防止训练中断时丢失进度。

恢复模型状态：通过加载.ckpt文件，你可以从之前保存的状态恢复训练，而不是从头开始。它包含了模型的权重、优化器的状态、学习率调度器的信息等。

攻击过程分析

攻击者通过在.ckpt文件中植入恶意代码，当通过pytorch框架中的torch.load()函数加载进模型，当使用pickle.load()函数反序列化操作时，就会执行恶意代码，从而修改optimizer优化器参数，从而影响模型训练结果。该攻击正是忽略了对.ckpt文件内容的检验，从而导致在正常训练过程中进行反序列化操作时，pickle.load()函数的不安全性导致了可以远程执行任意代码。CVE公布的是Transformers库4.38之前受到影响。以下是漏洞详细信息。

解决方案

避免使用pickle.load()处理不受信任的文件或数据，可以改用更安全的反序列化方法，如果必须使用，确保在隔离环境中执行，并严格限制数据来源。

验证数据来源，仅从受信任的来源加载模型检查点文件，或者对文件内容进行审查，确保来源可靠。

打入安全补丁，升级Hugging Face发布的最新版本。

加强代码审计，对所有人员提交的代码和数据进行严格审计，避免被植入恶意代码。

做好网络隔离，在本地部署大模型进行训练，避免直接与互联网接触。

重视安全管理，在大模型训练的各个环节加强安全管理，强化运维审计。

总结

AI安全的重要性不容忽视，尤其在训练阶段出现问题时，可能导致严重的后果和巨大的经济损失。一旦模型受到攻击或发生偏差，生成的结果可能误导用户，造成信誉损失和法律责任。最近某厂商的AI安全事件正是一个典型的例证，突显了在AI系统开发、部署、训练、测试等过程中，确保安全和可靠性的必要性。为了避免类似事件的发生，企业应加强对AI安全的重视，实施严格的测试和监控措施，以确保模型在各种环境下的稳健性和安全性。只有这样，才能在推动AI技术进步的同时，降低潜在的风险和损失。