tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Fix

Browse Source
This commit is contained in:
Dan Saunders
2025-09-22 16:10:41 -04:00
parent 92f2f6e73c
commit 9d69c6fb3e
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
scripts/benchmarks/deepseek_v3_moe.py
View File

@@ -193,8 +193,9 @@ def main() -> None: # pragma: no cover - CLI entrypoint
) )
def _uniform_gate(self, hidden_states): def _uniform_gate(self, hidden_states):
batch_tokens = hidden_states.shape[0] flat = hidden_states.view(-1, hidden_states.shape[-1])
return topk_idx[:batch_tokens], weights[:batch_tokens] token_count = flat.shape[0]
return topk_idx[:token_count], weights[:token_count]
patched_module.gate.forward = _uniform_gate.__get__( patched_module.gate.forward = _uniform_gate.__get__(
patched_module.gate, patched_module.gate.__class__ patched_module.gate, patched_module.gate.__class__