Add shared ModelParallel distribution-test helper by pctablet505 · Pull Request #2809 · keras-team/keras-hub

pctablet505 · 2026-07-17T09:03:38Z

Closes #2825

Part of #2807. First of a small series of per-model PRs; this one is just the shared foundation the rest build on.

Summary

Adds TestCase.run_distribution_test, a shared helper to be used by every model's test_distribution:

Sharding-spec assertions with a "matched at least one weight" check (kills dead assertions that silently pass against a regex matching nothing).
Full-weight coverage: every rank>=2 weight must be either explicitly mapped by the layout map or explicitly declared intentionally-replicated, or the test fails listing the offending paths.
A real forward pass, backward pass, and optimizer step — not just model construction.
A numerical-parity and short-training-convergence check against an undistributed twin model built from identical seeded weights (same input, same optimizer steps), at two mesh shapes chosen so at least one has >1 device on both the batch and model axes simultaneously.

Also:

Registers the multi_device pytest marker (previously used unregistered elsewhere in the repo, e.g. gpt2_causal_lm_test.py).
Adds a dedicated run_distribution_tests CI job that runs with 8 simulated CPU devices, scoped to that job's own environment only — the main run_tests job/matrix is untouched, so this doesn't change device visibility for the rest of the jax-backend test suite.

No model-specific changes in this PR — those land as separate small PRs on top of this one (see #2807's checklist).

Verification

ruff check / ruff format --check: clean.
YAML validated for the new CI job.
The helper itself is exercised end-to-end by the follow-up Gemma PR, which depends on this one.

Adds TestCase.run_distribution_test, a shared helper to be used by every model's test_distribution: sharding-spec + full-weight-coverage assertions, a real forward/backward/optimizer step (not just build()), and a numerical-parity + short-training-convergence check against an undistributed twin model. Registers the multi_device pytest marker (previously used unregistered elsewhere in the repo) and adds a dedicated CI job that runs distribution-related tests with 8 simulated CPU devices, scoped to that job only so the main test matrix's device visibility is unaffected.

gemini-code-assist

Code Review

This pull request registers a new multi_device pytest marker and introduces a robust run_distribution_test helper method in TestCase to standardize ModelParallel layout-map testing, verifying weight shardings, coverage, and numerical parity against undistributed twins. The review feedback highlights several edge cases where optional model outputs containing None values could cause AttributeError or conversion failures during flattening, structure mapping, and parity checks. Additionally, it suggests safely checking for the spec attribute on weight shardings to support non-NamedSharding weights, and optimizing the training convergence check by running fit once with epochs=5 instead of looping in Python.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Depends on keras-team#2809 (adds TestCase.run_distribution_test) for tests to pass -- this PR's own diff is independent (based directly on master, not stacked) so it stays small and reviewable on its own. Refactors Gemma's existing test_distribution/test_distribution_with_lora onto the shared helper. Adds test_layout_map_mesh_shapes (CI-safe sweep across representative 2D/3D mesh shapes) and test_layout_map_live_presets (manual, Kaggle-gated sweep across the full preset registry). Gemma's get_layout_map itself is unchanged -- already correct.

- Skip None leaves in the forward-output finiteness check and the zeros-like target construction, for backbones with optional outputs. - Fix parity_mesh_shapes DeviceMesh axis_names being hardcoded to a 2-tuple, which would raise on a 3D shape (dormant today since the default is 2D-only, but a real latent bug). - Compare parity outputs via assertAllClose directly on the nested structure instead of a manual tree.flatten()+zip(), which silently truncates on a structure-length mismatch instead of failing loudly. - Run the 5-step training-convergence check as a single fit(epochs=5) call instead of looping fit(epochs=1) 5 times -- same result, avoids 5x the per-call Keras/JAX setup overhead.

pctablet505 · 2026-07-17T09:50:32Z

@gemini-code-assist review

gemini-code-assist

Code Review

This pull request registers a new 'multi_device' pytest marker and introduces a comprehensive 'run_distribution_test' helper method in the TestCase class to run shared ModelParallel layout-map tests for backbones. The review feedback focuses on enhancing the robustness and correctness of this new test helper. Key recommendations include caching the CPU device list to avoid redundant calls, introducing a helper to safely copy 'init_kwargs' to prevent side-effects from in-place dictionary modifications, dynamically handling axis names for different mesh shape dimensions to avoid ValueErrors, and filtering out non-float tensors before checking for finite values to prevent backend errors.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

…ture - get_kwargs() helper: fresh dict copy each call so a constructor cannot mutate init_kwargs in place and leak into the next instantiation. - Fetch list_devices("CPU") once (all_devices) instead of re-querying on every parity mesh-shape iteration. - Skip non-float leaves (int/bool aux outputs) in the finiteness check via is_float_dtype -- np.isfinite on them is meaningless. - Build the compile() loss structure dynamically (None-preserving) to match the output structure instead of a bare "mse", which would break if any output leaf is None. - Guard allow_replicated against a bare string (iterates character by character instead of failing loudly).

…ame match The old -k "test_distribution or test_layout_map" selector matched by test-name substring repo-wide, so on this branch alone (no model files touched yet) it picked up pre-existing, not-yet-refactored tests in gemma/llama that predate this effort's device-pinning discipline and aren't this PR's responsibility -- e.g. IndivisibleError from an old hand-rolled mesh using all 8 simulated devices instead of a pinned 2. Switch to -m so only tests explicitly opted in via @pytest.mark.multi_device run here; verified locally this correctly narrows to the 3 pre-existing gpt2 multi_device tests on this branch (2 passed, 1 skipped) instead of also catching gemma/llama's unrelated failures. Tolerate exit code 5 (no tests collected) since the marker is only added by the per-model layout-map PRs, not this one -- expected on this branch standalone.

gemini-code-assist · 2026-07-17T12:39:18Z

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

- Forward-parity comparison: assert_same_structure first (fails loudly on a structure mismatch), then compare leaf-by-leaf skipping paired None leaves. A direct assertAllClose(undistributed_out, distributed_out) raises TypeError on a None leaf pair (tf.test.TestCase does None - None), so a backbone with an unset optional output would crash the parity check. - parity_mesh_shapes axis_names: handle 1D shapes explicitly. A 1D mesh shape previously got the 3-tuple (batch, seq, model) via the else branch, which DeviceMesh rejects (axis_names length must equal mesh rank).

test_distribution and test_layout_map_mesh_shapes both construct a real DeviceMesh/ModelParallel and need multiple simulated devices to mean anything, but neither carried the multi_device marker that PR keras-team#2809's dedicated CI job filters on. Without it, that job silently collects zero tests from this model. test_layout_map_live_presets already had the marker and stays excluded by design (kaggle_key_required + extra_large). For test_layout_map_mesh_shapes, the marker must go directly on the function, below @parameterized.named_parameters (not above it) -- absl.testing.parameterized's generated per-case methods only pick up pytestmark via functools.wraps when the mark decorates the original function itself; decorating the _ParameterizedTestIter object produced by named_parameters(...) does not propagate pytestmark to the expanded test cases. Verified empirically: outermost placement collected 0 of 12 parameterized cases under the marker filter, innermost placement collected all 12.

test_distribution, test_distribution_with_lora, and test_layout_map_mesh_shapes were not marked, so the dedicated multi_device CI job (added in keras-team#2809) was collecting zero real tests for this model -- only test_layout_map_live_presets carried the marker, and that test is deliberately excluded by the job's own kaggle_key_required/extra_large filter. For test_layout_map_mesh_shapes, which uses absl's parameterized.named_parameters, the marker must be applied as the innermost decorator (closest to the function) rather than outermost: named_parameters generates fresh method objects per expansion that don't inherit the pytestmark attribute pytest.mark sets on the original function, so an outer marker never reaches the expanded test cases. Verified empirically that only the inner-decorator ordering makes the parameterized cases collectible under -m multi_device.

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.

test_layout_map_live_presets was the only test in this file carrying @pytest.mark.multi_device, but it's also marked kaggle_key_required and extra_large, so PR keras-team#2809's CI job (which filters on "multi_device and not kaggle_key_required and not extra_large") was excluding it too -- leaving zero real tests collected from this model. Mark the two actual multi-device tests instead: test_distribution, and test_layout_map_mesh_shapes (the Tier-2 CI-safe mesh-shape sweep). For test_layout_map_mesh_shapes, the marker must go directly above the def, below @parameterized.named_parameters, not above it: absl's TestGeneratorMetaclass only forwards attributes set via functools.wraps inside _ParameterizedTestIter.__iter__ on the original function; a mark applied outermost lands on the _ParameterizedTestIter object itself and is silently dropped when the class dict is rewritten, so none of the 12 expanded per-mesh-shape test cases would inherit it. Verified via collect-only with the same "-m" filter CI uses: 13/19 tests selected (test_distribution + 12 mesh-shape cases), with test_layout_map_live_presets and the non-distribution tests correctly excluded either way.

test_distribution, test_distribution_vision, and test_layout_map_mesh_shapes build an actual ModelParallel/DeviceMesh and need multiple simulated devices, but were missing @pytest.mark.multi_device. Without it, the dedicated multi_device CI job (added by keras-team#2809) collects zero tests from this file -- only test_layout_map_live_presets carried the marker, and that test is deliberately excluded by the job's own "not kaggle_key_required and not extra_large" filter. For test_layout_map_mesh_shapes, the marker is placed as the innermost decorator (below @parameterized.named_parameters), not outermost: absl.testing.parameterized generates fresh function objects per expanded case and only copies pytestmark onto them when it wraps an already-marked function, not when the mark wraps the parameterized descriptor itself. Verified empirically that outermost placement causes pytest to collect 0 of the 18 expanded cases, while innermost placement collects all 18.

test_distribution, test_distribution_vision, test_distribution_audio, and test_layout_map_mesh_shapes each build a real ModelParallel DeviceMesh but were missing the multi_device marker, so the CI job added in keras-team#2809 (which filters on `-m "multi_device and not kaggle_key_required and not extra_large"`) was silently collecting zero tests from this file. test_layout_map_live_presets is left untouched -- it's already marked and correctly excluded from this CI job by the kaggle_key_required/extra_large filters. For test_layout_map_mesh_shapes, the marker had to go *below* @parameterized.named_parameters (innermost, closest to the function), not above it: absl's TestGeneratorMetaclass expands named_parameters by iterating the original undecorated test method, so a mark applied outside named_parameters never reaches the generated per-case test functions and silently fails to select any of them.

test_distribution and test_layout_map_mesh_shapes were not marked, so the dedicated multi_device CI job (added in keras-team#2809, filtering on -m "multi_device and not kaggle_key_required and not extra_large") would never collect any real test from this model, staying green while doing nothing. test_layout_map_live_presets already carries the marker (plus kaggle_key_required + extra_large) and is left as-is, correctly excluded from this job by design. For test_layout_map_mesh_shapes, the marker is placed as the innermost decorator (directly above the function, below @parameterized.named_parameters(...)), not outermost: absl.testing.parameterized expands each parameter case into a new function via functools.wraps, which only copies attributes already present on the wrapped function's __dict__ at expansion time. Placing pytest.mark above named_parameters sets pytestmark on the intermediate _ParameterizedTestIter object, which is discarded once the metaclass expands per-case methods, silently losing the marker on every generated case. Verified empirically via --collect-only: outer placement collected 0 of 12 parameterized cases under -m multi_device; inner placement collects all 12.

test_layout_map_live_presets was the only test in this file carrying @pytest.mark.multi_device, but it's also marked kaggle_key_required and extra_large, so PR keras-team#2809's CI job (which filters on "multi_device and not kaggle_key_required and not extra_large") was excluding it too -- leaving zero real tests collected from this model. Mark the two actual multi-device tests instead: test_distribution, and test_layout_map_mesh_shapes (the Tier-2 CI-safe mesh-shape sweep). For test_layout_map_mesh_shapes, the marker must go directly above the def, below @parameterized.named_parameters, not above it: absl's TestGeneratorMetaclass rewrites the class dict with per-case copies generated from the original function, and a mark applied outermost lands on the intermediate _ParameterizedTestIter object instead of the function, so it's silently dropped and none of the 18 expanded per-mesh-shape test cases would inherit it. Verified via collect-only with the same "-m" filter CI uses: 19/30 tests selected (test_distribution + 18 mesh-shape cases), with test_layout_map_live_presets and the non-distribution tests correctly excluded either way.

github-actions Bot requested review from divyashreepathihalli and laxmareddyp July 17, 2026 09:03

github-actions Bot added the Gemma Gemma model specific issues label Jul 17, 2026

pctablet505 mentioned this pull request Jul 17, 2026

Refactor Gemma distribution tests onto the shared helper #2808

Draft

gemini-code-assist Bot reviewed Jul 17, 2026

View reviewed changes

pctablet505 mentioned this pull request Jul 17, 2026

Tracking: per-model ModelParallel layout-map PRs (split of #2795), 17 PRs #2807

Open

pctablet505 mentioned this pull request Jul 17, 2026

DO NOT REVIEW: layout map work (integration/CI verification only) #2810

Draft

gemini-code-assist Bot reviewed Jul 17, 2026

View reviewed changes

Comment thread keras_hub/src/tests/test_case.py Outdated

Comment thread keras_hub/src/tests/test_case.py Outdated

Comment thread keras_hub/src/tests/test_case.py

Comment thread keras_hub/src/tests/test_case.py Outdated

Comment thread keras_hub/src/tests/test_case.py

pctablet505 marked this pull request as ready for review July 17, 2026 12:39

pctablet505 requested a review from jeffcarp July 22, 2026 11:02

Merge branch 'keras-team:master' into layout-map-test-helper

f1aea15

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add shared ModelParallel distribution-test helper#2809

Add shared ModelParallel distribution-test helper#2809
pctablet505 wants to merge 6 commits into
keras-team:masterfrom
pctablet505:layout-map-test-helper

pctablet505 commented Jul 17, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

pctablet505 commented Jul 17, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gemini-code-assist Bot commented Jul 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

pctablet505 commented Jul 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Verification

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

pctablet505 commented Jul 17, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gemini-code-assist Bot commented Jul 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

pctablet505 commented Jul 17, 2026 •

edited

Loading