From 216b09219e7b5d47c5f43c57a01898b065b80fa6 Mon Sep 17 00:00:00 2001 From: gulsumgudukbay Date: Thu, 3 Sep 2026 03:33:33 +0000 Subject: [PATCH] fix(mhc test): let the batch scale with the device count `_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. --- tests/unit/mhc_test.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/unit/mhc_test.py b/tests/unit/mhc_test.py index db0cd49c85..f05fe2ac2d 100644 --- a/tests/unit/mhc_test.py +++ b/tests/unit/mhc_test.py @@ -404,7 +404,6 @@ def test_use_mhc_pallas_kernel_dispatch(self, use_mhc_pallas_kernel): use_mhc_pallas_kernel=use_mhc_pallas_kernel, dim=128, sequence_length=256, - per_device_batch_size=1, dtype="bfloat16", ) with nn_partitioning.axis_rules(self.config.logical_axis_rules): @@ -481,7 +480,6 @@ def test_use_mhc_pallas_kernel_custom_block_size(self): mhc_pallas_kernel_bwd_feature_block_size=512, dim=128, sequence_length=128, - per_device_batch_size=1, dtype="bfloat16", ) with nn_partitioning.axis_rules(self.config.logical_axis_rules):