mirror of
https://github.com/exo-explore/exo.git
synced 2026-09-11 04:49:25 -04:00
qwen3_5_moe_split: drain — eval gathered instead of contribution
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
1 parent
95d4e66680
commit
312229a454
1 file changed
+1
-1
@@ -305,8 +305,8 @@ def pipelined_layer_loop(
|
||||
contribution = moe(last_layer, h_H1_pending)
|
||||
else:
|
||||
contribution = _zeros_like(h_H1_pending)
|
||||
mx.eval(contribution)
|
||||
gathered = mx.distributed.all_gather(contribution, group=group)
|
||||
mx.eval(gathered)
|
||||
out_H1 = gathered[MOE_RANK : MOE_RANK + 1]
|
||||
|
||||
# After the final B stage (Stage 2N-1), out_{N-1}_H0 = x_H0 (MOE's contribution
|
||||
|
||||
Reference in new issue
Block a user