sync

Jinxiaolong1129 · Nov 24, 2024 · b75aa65 · b75aa65
1 parent 9e5e6e9
commit b75aa65
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/collect-deepseek-predictor-data.py b/collect-deepseek-predictor-data.py
@@ -361,6 +361,7 @@ def _custom_ffn_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
     for name, module in model.named_modules():
         if type(module).__name__ == 'DeepSeekMoE':
             block_ffn_input_output_pair[name] = []
+            module._original_forward = module.forward
             module._module_name = name
             module.forward = _custom_ffn_forward.__get__(module, type(module))