Fixes in gemma3 tests and sharding

[vfdev-5]

Description:

• set flax from source in pyproject.toml while bonsai does not yet have a stable release
• removed ShardedLinear and ShardedEmbedding with nnx layers
• added sharding tests for gemma3 exectuted via selective tests on CPU, 4 emulated devices
• put a workaround to fix the issue with nnx.eval_shape when loading the model from pretrained.
• set "transformers>=4.30.0<5.0.0" in pyproject.toml

[gemini-code-assist]

Summary of Changes

Hello @vfdev-5, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly refactors the Gemma3 model to align with the latest flax.nnx sharding capabilities, moving away from custom sharding layers to utilize native nnx components. It includes critical workarounds for current nnx limitations related to model loading and attribute handling. The changes also encompass updates to the model's configuration for more granular sharding control, adjustments to the test suite to validate these sharding behaviors, and a dependency update for flax to ensure compatibility.

Highlights

• Flax NNX Sharding Migration: Replaced custom ShardedLinear and ShardedEmbedding implementations with native flax.nnx.Linear and flax.nnx.Embed layers, leveraging kernel_metadata and bias_metadata for explicit sharding configuration.
• Sharding Configuration Enhancements: Refined sharding configurations across VisionConfig, TextConfig, and ModelConfig classes, introducing new activation sharding specifications and simplifying the conditional application of sharding defaults.
• NNX Eval Shape Workarounds: Implemented workarounds for known nnx.eval_shape issues by dynamically computing position_ids and embed_scale within __call__ methods and using .set_value() for cache updates, preventing problems when loading pretrained models.
• Updated Flax Dependency: Modified pyproject.toml to fetch the flax library directly from its GitHub source, addressing compatibility needs before a stable release of bonsai.
• Enhanced Sharding Tests: Introduced new sharding tests with explicit multi-device setup (4 emulated CPU devices) and assertions for output sharding, ensuring correct distributed behavior of the Gemma3 model.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• bonsai/models/gemma3/modeling.py
  • Removed default_embed_init import and custom ShardedLinear and ShardedEmbedding classes.
  • Replaced custom sharded layers with flax.nnx.Linear and flax.nnx.Embed across various model components (e.g., VisionPatchEmbeddings, VisionAttention, SiglipMLP, Gemma3TextScaledWordEmbedding, Gemma3Attention, Gemma3MLP).
  • Introduced kernel_metadata and bias_metadata dictionaries for passing sharding configurations to nnx.Linear and nnx.Embed.
  • Added emb_patch_activation and emb_pos_activation to VisionShardConfig and its default method for more detailed activation sharding.
  • Refactored VisionConfig, TextConfig, and ModelConfig to simplify sharding configuration logic, defaulting to no_sharding() and applying specific sharding only when use_fsdp or use_tp is true.
  • Implemented workarounds for nnx.eval_shape by making position_ids and embed_scale dynamic and using .set_value() for cache updates.
  • Adjusted __call__ method signatures in Gemma3Model and forward function to explicitly use keyword arguments for cache, segment_ids, and token_type_ids.
• bonsai/models/gemma3/tests/test_outputs_gemma3.py
  • Updated the call to create_gemma3_from_pretrained to remove the mesh argument.
  • Modified the update of bonsai_model.embed_tokens.weight.embedding to use .set_value().
  • Commented out the assertion for embed_scale due to the nnx.eval_shape workaround.
  • Added .copy() to np.asarray(jp) when converting to torch.from_numpy to prevent potential issues.
  • Updated model call signatures to use keyword arguments for cache, segment_ids, and token_type_ids.
  • Adjusted the rtol parameter in np.testing.assert_allclose for test_full.
• bonsai/models/gemma3/tests/test_sharding_gemma3.py
  • Added @unittest.skipIf decorator to skip sharding tests if fewer than 4 JAX devices are available.
  • Changed the jax.make_mesh configuration from ((1, 1)) to ((2, 2)) to enable testing with more devices.
  • Switched from using gemma3_4b_it to a custom, smaller model configuration for sharding tests to ensure controlled testing.
  • Explicitly placed input arrays (n_img, n_text, n_tti, segment_ids) on devices with P(fsdp) sharding.
  • Added assertions to verify the output sharding of the model, checking for NamedSharding and P(fsdp, None, None) spec.
• pyproject.toml
  • Updated the flax dependency to point to its Git repository, allowing the use of the latest development version.
• scripts/run_selective_tests.py
  • Added a conditional block to run sharding-related tests (TestSharding) with JAX_NUM_CPU_DEVICES=4 and JAX_PLATFORMS=cpu environment variables, enabling emulated multi-device testing for sharding.

Activity

• The author, vfdev-5, initiated this pull request to address fixes in Gemma3 tests and sharing.
• Significant refactoring was performed to migrate the Gemma3 model's sharding implementation from custom layers to native flax.nnx components.
• Workarounds for nnx.eval_shape issues were introduced to ensure proper model loading and functionality.
• The flax dependency was updated to point to its source repository, indicating a need for a specific version or features not yet in a stable release.
• New sharding tests were added and existing tests were adjusted to validate the updated sharding mechanisms and model behavior in a multi-device environment.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces several significant improvements and fixes for the Gemma3 model implementation. The core change is the removal of custom ShardedLinear and ShardedEmbedding modules, replacing them with standard flax.nnx.Linear and flax.nnx.Embed layers, which simplifies the codebase. The PR also includes workarounds for issues with nnx.eval_shape when loading pretrained models. Additionally, new sharding tests have been added and the test execution script has been updated to run these tests on emulated CPU devices, which is a great addition for CI.

My review found one critical issue where incorrect sharding configurations are being used in SiglipMLP due to variable reuse. Apart from that, the changes are well-implemented and improve the maintainability and test coverage of the model.