Skip to content

Commit b587fc9

Browse files
committed
Mark Qwen distribution tests as multi_device for CI collection
test_distribution and test_layout_map_mesh_shapes exercise real DeviceMesh/ModelParallel sharding but were missing the multi_device marker, so the dedicated CI job (-m "multi_device and not kaggle_key_required and not extra_large") added in keras-team#2809 would collect zero tests from this file. For test_layout_map_mesh_shapes, the marker must go directly on the function, below @parameterized.named_parameters(...): absl's named_parameters expands each parameterization into a fresh bound function via functools.wraps(test_method), which only carries over attributes already present on the original function's __dict__, not attributes attached afterward to the _ParameterizedTestIter wrapper object that an outermost decorator would target. Placing the pytest mark outermost silently drops it from every expanded test case. test_layout_map_live_presets is already marked and intentionally excluded by this job's own kaggle_key_required/extra_large filters, so it is left untouched.
1 parent 7c2c235 commit b587fc9

1 file changed

Lines changed: 2 additions & 0 deletions

File tree

keras_hub/src/models/qwen/qwen_backbone_test.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -171,6 +171,7 @@ def test_num_parameters(self):
171171
model = QwenBackbone(**self.init_kwargs)
172172
self.assertEqual(model.count_params(), 1384)
173173

174+
@pytest.mark.multi_device
174175
def test_distribution(self):
175176
self.run_distribution_test(
176177
cls=QwenBackbone,
@@ -206,6 +207,7 @@ def test_distribution(self):
206207
for dims in (QWEN_0_5B_DIMS, QWEN_3B_DIMS, QWEN_7B_DIMS)
207208
for shape in CAPPED_MESH_SHAPES
208209
)
210+
@pytest.mark.multi_device
209211
def test_layout_map_mesh_shapes(self, dims, mesh_shape):
210212
if keras.backend.backend() != "jax":
211213
self.skipTest("`ModelParallel` testing requires the Jax backend.")

0 commit comments

Comments
 (0)