Repository navigation
Grouped expert dispatch: no padding to the busiest expert - #4
Merged
Merged
Conversation
Training the large Strix profile slowed to half speed as it went and then crashed. Measured on its checkpoint: routing grows "sink" experts that nearly every character picks among its top 8 - 1,960 of 2,081 characters at a typical row, 15x the fair share with 128 experts on the card. The dispatch padded every expert to the busiest one's load, so uncapped it computed ~16x the routed work and its buffers kept growing; capped, it dropped 30-60% of routing. - minagi/pool.py: the chosen experts are grouped by load (rounded with _bucket) and each group runs its own batched matmuls at its own size. Weights are gathered once per call and split into the groups' views; a gather per group was slower, because each one's backward allocates a gradient the size of the whole card. One character is a single group of top_k experts, so writing reads top_k experts' weights. The old single rectangle stays behind PooledMLP.grouped_dispatch = False. - Buffer capacities are rounded to a few dozen sizes (_bucket). ROCm on Windows ignores expandable_segments, and a new buffer size every call fragmented 17 GB of live tensors into 49 GB reserved. - train.py: allocated / reserved / peak GPU memory in the history file at every checkpoint. - tools/gpu_mem_sampler.ps1: per-process dedicated and shared GPU memory from Windows' performance counters. - Profiles and README: expandable_segments is ignored on Windows ROCm, and why capacity_factor stays 0 on the large profile. On the 8060S, step 1,938 of the large model, uncapped: training step 475 -> 813 char/s, peak 24.6 -> 18.8 GB, 0% dropped writing 5.3 -> 19.8 char/s A 45-minute run held 100-147 steps per 5 minutes (58 before) with reserved memory flat at 22.8 GB; held-out 2.17 -> 2.00. 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
Fixes the large Strix profile slowing to half speed during training and then crashing.
Cause. Routing grows "sink" experts that nearly every character picks among its top 8. On the trained checkpoint, the busiest expert received 1,960 of 2,081 characters at a typical row, 15× the fair share with 128 experts on the card. The dispatch padded every expert to the busiest one's load:
Fragmentation made it worse. ROCm on Windows ignores
expandable_segments, and a new buffer size on every call fragmented 17 GB of live tensors into 49 GB reserved.Changes
minagi/pool.py, grouped dispatch: the chosen experts are grouped by load and each group runs its own batched matmuls at its own size. Nothing is dropped thatcapacity_factordidn't drop, and the work is at most about 2× the routing's instead of n × the worst expert's.PooledMLP.grouped_dispatch = False.minagi/pool.py,_bucket: buffer capacities are rounded to a few dozen sizes (powers of two up to 128, then multiples of 128) so the allocator reuses blocks.train.py: allocated, reserved and peak GPU memory in the history file at every checkpoint.tools/gpu_mem_sampler.ps1: per-process dedicated and shared GPU memory from Windows' performance counters, sampled every 30 s.expandable_segmentsis ignored on Windows ROCm, and whycapacity_factorstays 0 on the large profile.Results (Radeon 8060S, Windows, ROCm 7.13 nightly, large model at step 1,938, uncapped)
45-minute run on this code:
The step rate eases while memory stays fixed, which fits the model learning to use more recurrent rows.
Not addressed here
With nearly the whole pool on the card, the balance term charges experts by admission share, which is close to uniform, so it can't see sink experts. Making it see routing load changes training behavior and needs its own test.
Test plan
pytest: 22/22.tests/test_dispatch.pychecks that the grouped path matches the rectangle's outputs and slot gradients, for one character and for an uneven-load chunk that splits into ≥ 2 groups, and that bucket padding changes nothing🤖 Generated with Claude Code