fix(mhc test): let the batch scale with the device count - #5133
Open
gulsumgudukbay wants to merge 1 commit into
Open
fix(mhc test): let the batch scale with the device count#5133gulsumgudukbay wants to merge 1 commit into
gulsumgudukbay wants to merge 1 commit into
Conversation
`_setup_mhc` uses `per_device_batch_size` as the *global* batch of `self.x`, which is why it defaults to `jax.device_count()`. These three tests overrode it to 1, so on any host with more than one device the batch is smaller than the mesh: the FSDP axis wants to partition axis 0 four ways on a 4-device runner and the layer dies with `IndivisibleError`. They pass on a single-device CPU, which is the only place CI runs them. Drop the override and take the default, as the rest of the file does; batch size is irrelevant to what these tests assert.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
gulsumgudukbay
marked this pull request as ready for review
September 4, 2026 04:43
gulsumgudukbay
requested review from
A9isha,
NuojCheng,
RissyRan,
SurbhiJainUSC,
abhinavclemson,
aireenmei,
bvandermoon,
darisoy,
dipannita08,
gagika,
gobbleturk,
hengtaoguo,
huytransformer,
igorts-git,
jiangjy1982,
khatwanimohit,
richjames0,
shralex,
shuningjin,
vipannalla and
xibinliu
as code owners
September 4, 2026 04:43
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
1 participant
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.
Description
_setup_mhcusesper_device_batch_sizeas the global batch ofself.x, which is why it defaults tojax.device_count(). Three tests override it to 1, so on any host with more than one device the batch is smaller than the mesh: the FSDP axis tries to partition axis 0 four ways on a 4-device host and the layer fails withIndivisibleError: array axis 0 is partitioned 4 times, but the dimension size is 1.They pass in CI only because the job that runs them (
cpu-unit) has a single device. These tests are auto-markedcpu_only, so the 4-GPUgpu-unitjob deselects them. Dropping the override and taking the helper's default keeps single-device behavior identical while making the batch divide the mesh anywhere else. Batch size is irrelevant to what these tests assert (kernel dispatch and block-size plumbing).Tests
Reproduce the failure on
mainand the fix on this branch with:Verified on 4x AMD Instinct MI355X (ROCm 10, jax 0.11.1): the three tests fail on
mainwith theIndivisibleErrorabove and pass on this branch. The fulltests/unit/mhc_test.pyfile also passes there(25 passed, 2 skipped).
Checklist
gemini-reviewlabel.