diff --git "a/2025-Ascend-Innovation-Contest/S1/MoE/\351\230\237\344\274\215emmm/README.md" "b/2025-Ascend-Innovation-Contest/S1/MoE/\351\230\237\344\274\215emmm/README.md" new file mode 100644 index 00000000..b29d1520 --- /dev/null +++ "b/2025-Ascend-Innovation-Contest/S1/MoE/\351\230\237\344\274\215emmm/README.md" @@ -0,0 +1,204 @@ +# README + +## 比赛内容(MOE赛道) + +MoE类： + +deepseek-ai/deepseek-moe-16b-chat + +Qwen/Qwen1.5-MoE-A2.7B-Chat + +在无精度误差的情况下提速这两个模型的prefill，decode和显存峰值 + +![img](./assets/score计算.png) + +## 最终成绩 + +![img](./assets/最终成绩.png) + +# 比赛复盘 + +## 前期思路 + +- flash-attention ： 由于确实是很常用的加速手段，原理上也work，所以基本上贯穿了优化策略，但实际上收益甚微，就连显存上也没有显示出优化 +- 算子融合 ： 一开始做的是诸如合并两次运算到一次运算里面；官方开会的时候提到了mindnlp.core.F，也提到了融合算子的下发调用开销可能和他的加速持平，实际测试下来没有加速，而且造成了精度误差(F.rms_norm) +- 合并python循环为矩阵运算 : 可以说是最有效的方法了，但是前期并没有探索的很深，浅浅掠过 +- 复用图/kernal ： 这个可谓是花了最多心思、同时又没弄出效果的方法，具体放在中期测试里面讲 +- 只遍历激活专家 ： 前期唯一的提分手段，从100->120 + +## 中期测试 + +- flash-attention + - 通过简单网络来测试，flash-attention对于长序列确而有提速效果，但是在中短序列不明显，有时候还会因为未知波动效果不如baseline + - 官方接口 `mindspore.ops.flash_attention_score`会带来一定的精度误差，具体而言qwen的prompt2会mismatch +- 算子融合 + - F.rms_norm 不仅没加速还带来了精度误差(应该是qwen的prompt1会mismatch)，遂直接放弃；对于review中提到的融合算子精度对齐没有缺陷，我猜测可能是进入F.rms_norm前所必须做的精度转化操作导致的，虽然我当时尝试了float32也还是有mismatch + - 但是我没太理解会议里面讲的要比较下放损耗和融合算子加速效果，我个人仍然觉得这应该要work，但是却没有 +- Graph&Pynative mode - kernal/图复用 + - 一开始打算用分桶填充策略，设置 `seq_len = [1,2,4,8,..,128]`的桶来多次调用模型生成来生成这些尺寸的图，为输入的prompt寻找恰好不小于他的桶进行padding触发图复用，但是毫无效果，于是开始探索图复用的条件，网上有说法是需要 `@mindspore.jit`即时编译/`Graph mode`静态图模式才能生成可以复用的图，于是进入下一步测试 + - @mindspore.jit：基本用不了，对于网络模型底层的try-except控制流，jit不支持这种低效分支，而我们又不大可能去修改底层的控制流(太多)，遂放弃 + - Graph mode：经过我用简单网络的多次测试(充分预热，多次测试取平均)，Graph模式比Pynative模式慢10倍左右，匪夷所思，初步怀疑是没触发图复用所以编译开销也算里面了(也就是单纯用Graph跑一次不会建立图？要显性调用什么函数或者装饰器吗)，遂半放弃 + - static-cache ：没做成功，因为需要把动态cache 换成 static cache，bug较多，时间上不允许，而且直播的时候说提升不大。 +- Profiler + - 这是一个很好的工具(疑似)，但是直到最后都不知道如何使用,一方面是断点设置和信息收集的问题，但这个问题不大 + - ![img](./assets/mindstudio.png) + - 最重要的是这个页面我只看到NPU的free/compute比值很大，除此之外不知道如何分析来调优了,要是能看**别人实际调优一遍肯定会好很多，求教程！！** +- MOE分析 + - 通过模型原来的代码，在self.mlp = ... 这一行，我发现了有一个if控制流，走moe/mlp,尝试使用走mlp之后，prefill/decode耗时降低了**20倍**，这时候我才意识到，原来前面有的没的都是**次要矛盾**，只要把**moe这个模块的代码**优化好了，就已经胜利了 + - 测试对比了moe和attention模块的占时，发现moe耗时是attention的20倍左右，所以压根不需要管attention部分的优化，也说明了为什么flash-attention部分没什么提速效果，因为**attention对总时长的影响实在微乎其微。** + - 通过查看榜单，思考前两名prefill=1700的同时为什么显存都有略微下降，可以明确的是他们肯定用空间换了时间，但是具体是做了什么呢？ + - 通过对moe模块的分析，终于才把握住了主要矛盾；结合前面总总，不难推断肯定是moe部分把某些**串行换并行或者大矩阵计算**，从显存情况我更倾向于是**大矩阵计算**，所以问题就转移到如何做这个大矩阵运算 + +## 后期优化 + +以下代码是我们moe prefill 和 decode 最快的一版 + +### prefill + +1. **Pad (填充)**: 将分配给不同专家的、数量不等的“锯齿状”token数据，通过 tensor_scatter_update 填充成一个规整的、[专家数, 最大Token数, 隐藏层大小] 的“矩形”张量。 +2. **BMM (批量矩阵乘法)**: 利用这个规整的张量，调用一次 ops.bmm 即可同时计算所有专家的输出，将硬件并行度拉满。 +3. **Gather (收集)**: 计算完成后，用 gather_nd 从填充后的结果中高效地提取出有效的输出数据。 + +```Python + @no_grad() + def moe_infer_prefill_fast(self, x, flat_expert_indices, flat_expert_weights): + num_total_assignments = flat_expert_indices.shape[0] + hidden_size = x.shape[-1] + num_experts = len(self.experts) + + # 1) 排序专家分配 + idxs = flat_expert_indices.argsort() + sorted_expert_indices = flat_expert_indices[idxs] + sorted_token_indices = idxs // self.num_experts_per_tok + permuted_tokens = x[sorted_token_indices] + sorted_weights = flat_expert_weights[idxs] + + # 2) 计算Padding所需的尺寸 + tokens_per_expert = sorted_expert_indices.bincount(minlength=num_experts) + max_tokens_per_expert = tokens_per_expert.max().item() + + if max_tokens_per_expert == 0: + return ops.zeros_like(x) + + # 3) 创建Padded张量 + expert_offsets = ops.cumsum(tokens_per_expert, dim=0) - tokens_per_expert + token_indices_in_sorted = mnp.arange(num_total_assignments) + relative_pos_in_expert = token_indices_in_sorted - expert_offsets[sorted_expert_indices] + + gather_indices_sparse = ops.stack([sorted_expert_indices, relative_pos_in_expert], dim=1) + + # --- 关键修正: 使用 tensor_scatter_update --- + padded_tokens = ops.zeros((num_experts, max_tokens_per_expert, hidden_size), dtype=x.dtype) + # 直接使用 mindspore.ops.tensor_scatter_update + padded_tokens = mindspore.ops.tensor_scatter_update(padded_tokens, gather_indices_sparse, permuted_tokens) + + # 4) 堆叠所有专家权重 + gate_weights = ops.stack([expert.gate_proj.weight for expert in self.experts], dim=0) + up_weights = ops.stack([expert.up_proj.weight for expert in self.experts], dim=0) + down_weights = ops.stack([expert.down_proj.weight for expert in self.experts], dim=0) + + # 5) --- 核心：巨型批处理矩阵乘法 (BMM) --- + gate_out = ops.bmm(padded_tokens, gate_weights.transpose(0, 2, 1)) + up_out = ops.bmm(padded_tokens, up_weights.transpose(0, 2, 1)) + act_out = self.experts[0].act_fn(gate_out) * up_out + padded_expert_outputs = ops.bmm(act_out, down_weights.transpose(0, 2, 1)) + + # 6) 从Padded结果张量中按原始顺序gather回所有有效结果 + # mindspore.ops.gather_nd 是 tensor_scatter_update 的逆操作，非常适合此场景 + expert_outputs_sorted = mindspore.ops.gather_nd(padded_expert_outputs, gather_indices_sparse) + + # 7) 最终加权和还原 + final_output = ops.zeros_like(x) + final_output = mindspore.mint.scatter_add( + final_output, + 0, + sorted_token_indices.view(-1, 1).tile((1, hidden_size)), + expert_outputs_sorted * sorted_weights + ) + return final_output +``` + +### Decode + +1. **向量化计算** + +新代码先将这个 token 选出的 top_k 个专家的weight收集起来，用 ops.stack 堆叠成一个批次。然后，通过 **ops.bmm** 这个算子，一次性并行完成这topk个专家的所有计算。 + +1. **内存局部性优化** + +通过 init_active_expert_cache 函数提前进行预处理。在模型预热（阶段识别出最常用的专家，然后将这些“热门专家”的权重从它们各自零散的位置提前抽调出来，用 ops.stack 堆叠成一个巨大且连续的内存块（即 self.cache_gate_w 等缓存张量）。在实际解码时，代码会优先从这个连续的缓存块中通过索引（self.cache_gate_w[eid]）直接读取权重。这种直接在大块连续内存上的索引操作速度极快，因为它能高效利用硬件的内存缓存机制（“缓存命中” Cache Hit），避免了复杂的对象查找。 + +```Python + def init_active_expert_cache(self, active_ids): + """ + 预热后调用，将常用专家的权重预先提取并堆叠， + 形成一个内存连续的“快速访问缓存”。 + """ + self.cache_gate_w = ops.stack([self.experts[i].gate_proj.weight for i in active_ids], dim=0) + self.cache_up_w = ops.stack([self.experts[i].up_proj.weight for i in active_ids], dim=0) + self.cache_down_w = ops.stack([self.experts[i].down_proj.weight for i in active_ids], dim=0) + + def moe_infer_decode_fast(self, x, flat_expert_indices, flat_expert_weights): + """ + 利用“权重缓存”和“BMM向量化”实现极致解码速度。 + """ + top_k = flat_expert_indices.shape[0] + hidden_size = x.shape[-1] + + selected_gate_w = [] + selected_up_w = [] + selected_down_w = [] + + # 1. 核心：从“快速缓存”或“慢速原始列表”中收集权重 + for eid in flat_expert_indices.tolist(): + # 检查缓存是否存在且eid在缓存范围内，如果满足则进入“快速通道” + if hasattr(self, "cache_gate_w") and eid < self.cache_gate_w.shape[0]: + selected_gate_w.append(self.cache_gate_w[eid]) + selected_up_w.append(self.cache_up_w[eid]) + selected_down_w.append(self.cache_down_w[eid]) + else: # 否则，回退到“慢速通道” + selected_gate_w.append(self.experts[eid].gate_proj.weight) + selected_up_w.append(self.experts[eid].up_proj.weight) + selected_down_w.append(self.experts[eid].down_proj.weight) + + # 2. 将收集到的分散权重堆叠成一个批次 + selected_gate_w = ops.stack(selected_gate_w, dim=0) + selected_up_w = ops.stack(selected_up_w, dim=0) + selected_down_w = ops.stack(selected_down_w, dim=0) + + # 3. 向量化计算：使用BMM一次性完成所有专家运算 + x_expanded = x.expand((top_k, 1, hidden_size)) + gate_out = ops.bmm(x_expanded, selected_gate_w.transpose(0, 2, 1)) + up_out = ops.bmm(x_expanded, selected_up_w.transpose(0, 2, 1)) + intermediate_states = self.experts[0].act_fn(gate_out) * up_out + expert_outputs = ops.bmm(intermediate_states, selected_down_w.transpose(0, 2, 1)) + + # 4. 向量化聚合 + weighted_sum = (expert_outputs * flat_expert_weights.unsqueeze(-1)).sum(axis=0) + return weighted_sum +``` + +#### 其他trick + +- trick1 ： **劫持预热**，事先过一遍短中长三种prompt的generate，充分预热，当时的想法是瞎猫碰一下死耗子看看能不能触发图复用，结果意外确实降低了prefill时延，由于没有充分实验不确定原理 +- trick2 ： 根据测试的三个prompt的长度，**用Prompt=0/1/2控制走的优化流**，部分优化流对某些Prompt会带来精度误差，这样是为了在实在解决不了精度问题的情况下,不至于直接放弃这个有效的优化，利用其他Prompt的优化先快带动后快 +- trick3 ： **init_active_expert_cache和warmup_moe_model_deep** + - 在预热的时候，记录下所有被激活过的专家的ID，缓存那些在预热中被激活过的active_ids的权重（ops.stack）。 + - 如果缓存已经建立，并且当前需要的专家 eid 就在缓存里，它会直接从连续的 cache_gate_w 张量中索引权重。 + + + +## 收益点 + +| 策略名称 | 说明 | 显存峰值 | Prefill | Decode | 总分 | +| :----------------------------------------------: | :----------------------------------------------------------- | :-------------: | :---------------: | :---------------: | :---------------: | +| DeepseekMoe + Qwen MoE模块优化 | Decode直接遍历激活专家 | 100→100 | 100→132 | 100→400 | 100→200 | +| Rotary优化 | 用`ops.split`替代`rotate_half`切片方式 | 100→100 | 133.4445→132.4821 | 427.7311→437.5848 | 220.919→223.3556 | +| moe_prefill_fast / moe_decode_fast | 串行专家计算改为大批量并行BMM，减少Python循环，速度更快（LongPrompt dispatch） | 100→98.4848 | 132.4821→163.8114 | 437.5848→454.7424 | 223.3556→239.0129 | +| init_active_expert_cache / warmup_moe_model_deep | 缓存预热期间激活专家权重，直接索引cache提升性能 | 98.4848→98.4848 | 163.8114→198.4985 | 454.7424→493.2538 | 239.0129→263.4124 | +| Pad→BMM→Gather流程 | 将专家计算合并为一次BMM，保证精度float32并按LongPrompt dispatch | 98.4848→83.3333 | 198.4985→487.1616 | 493.2538→490.5996 | 263.4124→353.6982 | + +## 总结 + +- 对于精度优化比赛，不应该一上来就花费大量时间在框架本身优化、常用优化等上，应该先通过充分测试找到**主要矛盾**，因为一般这种比赛都会有侧重点，比如这个比赛就是moe部分，如果第一天我就能测试出moe的时间占比如此浮夸，我想我就不会把时间放在细枝末节上面 +- 对于调试工具(如Profiler等)，这是辅助完成第一步的，很有学习的必要，这种可视化工具分析的能力或者输出、断点分析能力将是打比赛的重要能力，说是最重要也不为过，这样才能有的放矢 \ No newline at end of file diff --git "a/2025-Ascend-Innovation-Contest/S1/MoE/\351\230\237\344\274\215emmm/assets/mindstudio.png" "b/2025-Ascend-Innovation-Contest/S1/MoE/\351\230\237\344\274\215emmm/assets/mindstudio.png" new file mode 100644 index 00000000..a963d5c6 Binary files /dev/null and "b/2025-Ascend-Innovation-Contest/S1/MoE/\351\230\237\344\274\215emmm/assets/mindstudio.png" differ diff --git "a/2025-Ascend-Innovation-Contest/S1/MoE/\351\230\237\344\274\215emmm/assets/score\350\256\241\347\256\227.png" "b/2025-Ascend-Innovation-Contest/S1/MoE/\351\230\237\344\274\215emmm/assets/score\350\256\241\347\256\227.png" new file mode 100644 index 00000000..bbae9015 Binary files /dev/null and "b/2025-Ascend-Innovation-Contest/S1/MoE/\351\230\237\344\274\215emmm/assets/score\350\256\241\347\256\227.png" differ diff --git "a/2025-Ascend-Innovation-Contest/S1/MoE/\351\230\237\344\274\215emmm/assets/\346\216\222\350\241\214\346\246\234.png" "b/2025-Ascend-Innovation-Contest/S1/MoE/\351\230\237\344\274\215emmm/assets/\346\216\222\350\241\214\346\246\234.png" new file mode 100644 index 00000000..00a37c31 Binary files /dev/null and "b/2025-Ascend-Innovation-Contest/S1/MoE/\351\230\237\344\274\215emmm/assets/\346\216\222\350\241\214\346\246\234.png" differ diff --git "a/2025-Ascend-Innovation-Contest/S1/MoE/\351\230\237\344\274\215emmm/assets/\346\234\200\347\273\210\346\210\220\347\273\251.png" "b/2025-Ascend-Innovation-Contest/S1/MoE/\351\230\237\344\274\215emmm/assets/\346\234\200\347\273\210\346\210\220\347\273\251.png" new file mode 100644 index 00000000..9e5b9d8e Binary files /dev/null and "b/2025-Ascend-Innovation-Contest/S1/MoE/\351\230\237\344\274\215emmm/assets/\346\234\200\347\273\210\346\210\220\347\273\251.png" differ diff --git "a/2025-Ascend-Innovation-Contest/S1/MoE/\351\230\237\344\274\215emmm/patches.zip" "b/2025-Ascend-Innovation-Contest/S1/MoE/\351\230\237\344\274\215emmm/patches.zip" new file mode 100644 index 00000000..6a7b63a5 Binary files /dev/null and "b/2025-Ascend-Innovation-Contest/S1/MoE/\351\230\237\344\274\215emmm/patches.zip" differ