Fix LoRA fine-tuning for Gemma4 models

[copybara-service]

Fix LoRA fine-tuning for Gemma4 models

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[darknecrocities]

This is a substantial and well-executed change that improves checkpoint robustness and significantly strengthens the model’s ability to handle real-world structural drift across LoRA, multimodal extensions, and nn.share_scope variations. The introduction of _needs_reconciliation and _reconcile_tree is particularly valuable, as it cleanly isolates structural repair logic without impacting the default Gemma3 / legacy restore path.

The multimodal parameter handling is also more systematic now, with consistent support for both vision and audio encoders and their associated projection/norm keys. Consolidating these into _MM_TOP_LEVEL_KEYS and _MM_EMBEDDER_KEYS improves maintainability and reduces the risk of silent parameter drops when new modalities are introduced.

The added LoRA support in _SUPPORTED_MODULES and the expanded test coverage are strong improvements, especially given the complexity introduced by multiple Einsum variants across Gemma3n, Gemma4, and nano layers. The reconciliation tests are particularly thorough and effectively validate both stub removal and leaf-vs-dict normalization behaviors.

A few minor considerations: _needs_reconciliation performs recursive structural checks that may become expensive for very large parameter trees; if this path is hit frequently, it may be worth benchmarking or caching intermediate structural signatures. Additionally, the reconciliation logic assumes that empty dicts always represent LoRA stubs, which is correct for current usage but may benefit from a more explicit tagging mechanism in the future to avoid accidental false positives if other empty scopes are introduced.