Gemma already had test_distribution/test_distribution_with_lora, but they predate the shared helper from #2825, so this rewrites them onto run_distribution_test. While in there, get_layout_map needed a real fix: it was sharding the key/value kernels on the model axis, but those are sized by num_key_value_heads, which MQA/GQA configs can drop to 1, so any mesh whose model dimension didn't divide it raised IndivisibleError. Leaving k/v replicated fixes that; the query rule is untouched. Also adds a mesh-shape sweep across realistic pod-like topologies and a Kaggle-gated live-preset check with a memory budget.
One of the layout-map tasks under #2807. PR #2808.
Gemma already had
test_distribution/test_distribution_with_lora, but they predate the shared helper from #2825, so this rewrites them ontorun_distribution_test. While in there,get_layout_mapneeded a real fix: it was sharding the key/value kernels on the model axis, but those are sized bynum_key_value_heads, which MQA/GQA configs can drop to 1, so any mesh whose model dimension didn't divide it raisedIndivisibleError. Leaving k/v replicated fixes that; the query rule is untouched. Also adds a mesh-shape sweep across realistic pod-like topologies and a Kaggle-gated live-preset check with a memory budget.One of the layout-map tasks under #2807. PR #2808.