MoEfication代码实践

MoEfication的代码:https://github.com/thunlp/MoEfication

如何跑通

  • 小修改就能跑通的代码在examples文件夹下,下面介绍使用K-means聚类划分专家+训练MLP构建路由模块

  • t5_cluster_example.py 这个文件使用k-means聚类构建专家。这个程序会生成结尾是wi.weight的文件,它表示每一层的FFN中的每一个神经元(也就是矩阵的每一列)所属的专家

  • 训练MLP涉及到两个程序:

    • t5-sst2-inf.py:它使用一个情感分析的数据集sst-2,让t5-base执行推理,并将每一个层的推理结果都保存下来,所以非常消耗存储。将完整训练数据的推理结果保存下来,需要大约50G

    • t5_select_example.py:训练MLP,也就是根据上面保存过的推理结果,结合构建的专家,来训练路由模块

  • 测试MLP构建效果:t5-sst2-mlp.py

注意事项

  • 在这个代码中,对T5-base进行了拆分。由于T5是encoder-decoder结构的模型,在拆分中将encoder和decoder分别进行拆分。T5-base有12层,所以要拆分24次。

  • 如果要对当前的decoder-only结构的模型进行拆分,需要依据实际模型结构、层数改代码。

  • 而且,这个代码在前向传播过程中,依然是把整个FFN加载进来,然后根据聚类得到的每个神经元属于哪个专家来做pattern,这个过程仍然是会将整个FFN从HBM移动到SRAM中。由于这是访存瓶颈的,所以减少的计算量并不会减少整体时间。如果想要减少访存,需要在得到每个神经元所属专家后,将它们提取出来单独保存。

训练MLP的过程

  • 首先让完整的t5-base对训练数据集sst-2进行推理,每条训练数据的推理结果是一个seq_len*d_model大小的张量,然后对所有推理结果进行concat并保存。

  • 在训练过程中,对于concat后的数据,每次取512个token,也就是512*d_model大小的tensor。

  • 将这个tensor先在完整FFN中获得激活值,然后根据包含了神经元所属哪个专家信息的pattern来进行mask,得到的值作为label;

  • 然后将tensor经过MLP处理得到pred,计算pred和label的损失,进行反向传播

训练的部分代码:

train_hiddens = train_hiddens[torch.randperm(train_hiddens.size()[0])]  # 打乱顺序

            train_pbar = tqdm.tqdm(range(0, train_hiddens.shape[0], 512))
            for i in train_pbar:
                model.zero_grad()

                input = train_hiddens[i:i + 512, :].float().clone().detach().cuda()
                with torch.no_grad():
                    acts = torch.relu((torch.matmul(input, ffn_weight))).float()  # 所有神经元对这些token的预测
                    scores = torch.matmul(acts, patterns)  # 表示每个专家对这些token的预测
                    scores /= scores.max()  # 得到每个专家的得分
                pred = model(input)  # 得到MLP对每个专家的得分预测
                loss = loss_func(pred.view(-1), scores.view(-1))

                loss.backward()
                optim.step()

                train_pbar.set_description("loss: {:.4f}".format(loss.item()))

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值