Refactor Gemma distribution tests onto the shared helper

[pctablet505]

Closes #2826

Part of #2807. Depends on #2809 (adds TestCase.run_distribution_test, used by the tests below) for CI to pass — this PR's own diff is independent of #2809 (based directly on master, not stacked on its branch), so it stays small and reviewable on its own regardless of merge order. CI on this PR will go green once #2809 merges.

Summary

Refactors Gemma's existing test_distribution/test_distribution_with_lora onto the shared helper from #2809. get_layout_map is fixed to leave key/value kernels replicated on the model axis (they are sized by num_key_value_heads, which MQA/GQA configs make as small as 1, so sharding them raised IndivisibleError whenever the model mesh dimension did not divide it); the query rule is unchanged.

Also adds:

• test_layout_map_mesh_shapes: a sweep across representative 2-D and 3-D device-mesh shapes sized to mirror real accelerator-pod topologies, using realistic (not toy) dimensions frozen as literals, with an explicit skip (not silent drop, not failure) for mesh shapes where num_query_heads doesn't divide the model-parallel axis size — an inherent tensor-parallelism ceiling, not a bug.
• test_layout_map_live_presets: a manual, Kaggle-credential-gated sweep against every real preset in the registry (marked kaggle_key_required + multi_device + extra_large so it never runs in standard CI), with a per-width-class memory-budget estimate that skips (with a logged reason) any config too large to safely build locally.

Verification

Verified against a local checkout that already includes #2809's helper (since this PR's own base doesn't have it yet):

• keras_hub/src/models/gemma/ full suite: 45 passed, 22 skipped, KERAS_BACKEND=jax.
• test_distribution + test_distribution_with_lora: pass, including the parity check at (1,2) and (2,2) meshes.
• test_layout_map_mesh_shapes: 18/18 pass across 3 representative Gemma-family width classes × 6 mesh shapes (2x4, 1x8, 4x4, 2x2x2, 1x1x8, 2x2x4).
• test_layout_map_live_presets: fetches and dedupes all 23 registry presets into width classes, validates divisibility across the full mesh-shape matrix for each.
• ruff check and ruff format --check: clean.

Limitations

The model-parallel mesh axis can't usefully exceed a preset's num_query_heads, since that's the axis actually being tensor-parallelized — see the design doc's "hard limit" section for why.

For Gemma specifically, the safe ceiling depends on which preset you're sharding:

preset	num_query_heads	safe model-axis up to	first failing model-axis
c2s_scale_gemma_2_27b_en	32	32	64
c2s_scale_gemma_2_2b_en	8	8	16

The two presets above are Gemma-architecture models hosted under other Kaggle organizations (cell2sentence); the rest of the registry's gemma presets sit behind Kaggle's consent gate and couldn't be pulled to confirm numbers, so they're omitted here rather than guessed at.

[gemini-code-assist]

Code Review

This pull request refactors and enhances the distribution testing framework for KerasHub. It introduces a reusable run_distribution_test helper in TestCase to standardize model-parallel layout-map testing, including forward/backward passes and numerical parity checks against undistributed twins. It also updates the Gemma backbone's layout map to prevent sharding indivisibility errors on key/value kernels and adds extensive multi-device tests for Gemma. The review feedback focuses on improving the robustness and efficiency of the new test helper—specifically by dynamically handling 3D mesh shapes, avoiding fragile manual output flattening, and optimizing the training convergence loop—as well as ensuring proper resource management when opening configuration files.

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.

[pctablet505]

@gemini-code-assist review

[gemini-code-assist]

Code Review

This pull request updates the Gemma backbone layout map to replicate key/value kernels on the model-parallel axis while keeping query kernels sharded, preventing potential division errors during model parallelism. It also adds extensive tests to verify layout maps across various mesh shapes and live presets. The review feedback points out critical test setup issues: the test class must be decorated with @parameterized.use_parameterized to run the parameterized tests, and Gemma 2 specific architecture flags should be preserved in both the parameterized dimensions and the live preset test configurations to avoid testing Gemma 2 models under the Gemma 1 architecture.

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.

[gemini-code-assist]

The GemmaBackboneTest class uses @parameterized.named_parameters on test_layout_map_mesh_shapes, but the class itself is not decorated with @parameterized.use_parameterized and does not inherit from parameterized.TestCase. Without this decorator, absl.testing.parameterized will not expand the parameterized test cases, causing them to be ignored or fail during test collection. Please decorate the class with @parameterized.use_parameterized to ensure the parameterized tests are executed correctly.

@parameterized.use_parameterized
class GemmaBackboneTest(TestCase):

References

1. Poke Holes in the Implementation: Your goal is to critically test the logic. Actively search for and point out failing edge cases, race conditions, or unhandled exceptions in the implementation.

[pctablet505]

No change needed — TestCase (in test_case.py) already inherits parameterized.TestCase (class TestCase(tf.test.TestCase, parameterized.TestCase)), so GemmaBackboneTest(TestCase) gets it transitively. @parameterized.named_parameters does expand correctly (verified locally, 6 separate test cases collected for test_layout_map_mesh_shapes).

[gemini-code-assist]

The GEMMA2_9B_DIMS dictionary is missing Gemma 2 specific architecture flags like use_post_attention_norm and use_post_ffw_norm. Without these flags, the parameterized test test_layout_map_mesh_shapes will build and test Gemma 2 dimensions using the Gemma 1 architecture. Please add these flags to make the test realistic and robust.

Suggested change

	GEMMA2_9B_DIMS = {
	"source_preset": "gemma2_9b_en (real ratio, memory-scaled dims)",
	"vocabulary_size": 2048,
	"num_layers": 1,
	"num_query_heads": 32,
	"num_key_value_heads": 16, # real ratio: GQA, 2:1.
	"hidden_dim": 512,
	"intermediate_dim": 2048,
	"head_dim": 32,
	}
	GEMMA2_9B_DIMS = {
	"source_preset": "gemma2_9b_en (real ratio, memory-scaled dims)",
	"vocabulary_size": 2048,
	"num_layers": 1,
	"num_query_heads": 32,
	"num_key_value_heads": 16, # real ratio: GQA, 2:1.
	"hidden_dim": 512,
	"intermediate_dim": 2048,
	"head_dim": 32,
	"use_post_attention_norm": True,
	"use_post_ffw_norm": True,
	}

References

1. Demand Robustness: Do not accept fragile code. If the proposed code is not robust enough or lacks proper error handling, explicitly tell the author why the current approach is brittle and what must be done to reinforce it.