Repository navigation
Faster writing: run only the chosen experts; uncap the large profile - #3
Merged
Merged
Conversation
Writing a character was memory-bandwidth bound. Each character is a forward of its own and routes to top_k experts, but the dispatch ran its batched matmuls - and cast fp32 weights to bf16 - over every expert on the card. With 128 on the card that read 604M weights per row to use 8. - pool.py: when at most half the card's experts were chosen, gather just those before the matmuls. Training chunks reach nearly every expert, so they keep the old path (A/B on a training step: 4,794 vs 4,798 ms). tests/test_dispatch.py checks outputs and slot gradients match. - Fewer per-row host syncs and tiny launches (134 -> 67 syncs a character): cached slot index tensor in PagedPool, last_seen read once per _place, telemetry EMAs kept on the device, cap known for a single character, admit() indexes a list instead of a tensor per expert, one write to `ever` per admission. Writing speed on the 8060S (tools/bench_hw.py --weights): large model (128 on the card) 2.4 -> 9.4 char/s trained model (32 on the card) 7.8 -> 11.0 char/s - configs/strix_halo_large.yaml: capacity_factor 1.5 -> 0. At 1.5 a fresh model dropped 59% of routing (76% in a real run); uncapped a step is 27% slower and peaks at 22.8 GB instead of 17.8. - tools/bench_hw.py: a writing benchmark (--weights to use a real directory) and the dropped-routing share on the training-step line. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Writing is much faster, and the large profile no longer drops most of its routing.
Why writing was slow. Each character the model writes is a forward pass of its own, and it routes to
top_k= 8 experts. The dispatch, though, ran its batched matmuls over every expert on the card, converting each one's fp32 weights to bf16 on every call. With 128 on the card (strix_halo_large.yaml), every recurrent row read 604M weights in order to use 8 experts. Writing was limited by memory bandwidth.Changes
minagi/pool.py: when at most half the card's experts were chosen, the dispatch gathers just those weights before the matmuls.Fewer per-row host syncs and tiny launches, 134 → 67 syncs per character:
PagedPool, instead of a host→device copy on every routing calllast_seenread once per_place, instead of one sync per slotpressure,want_k) kept on the deviceadmit()indexes a Python list rather than a tensor per experteverper admissionconfigs/strix_halo_large.yaml:capacity_factor1.5 → 0 (uncapped). A fresh model with 128 experts on the card piles onto a few of them.A real run's first minutes dropped 75.6% at 1.5.
tools/bench_hw.py: a writing benchmark (--weights DIRuses a real directory, read-only) and the dropped-routing share on the training-step line.Results (Radeon 8060S, Windows, ROCm 7.13 nightly)
weights-strix, 128 on the card)weights, 32 on the card)What remains is fixed per-row overhead: about 150 small kernels plus the halting syncs. The untrained large model still runs nearly all 24 rows per character. HIP graphs or
torch.compilewould be the next step.Test plan
pytest: 19/19, including the newtests/test_dispatch.py. It checks that the subset dispatch gives the same outputs and slot gradients as the full one, and that the subset path is actually taken for a single character🤖 Generated with Claude Code