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()))