Code updates in qwen3

[vfdev-5]

• no sharding like in gemma3
• replaced custom rmsnorm, einsum classes with nnx
• sharding test

[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 Qwen3 model's sharding mechanisms and integrates more deeply with Flax NNX's native capabilities. The changes aim to provide a more robust and flexible distributed training setup by replacing custom sharding logic and module implementations with standard Flax components. This not only simplifies the codebase but also improves the model's compatibility with advanced sharding strategies like FSDP and TP, ensuring better performance and scalability on multi-device setups.

Highlights

• Sharding Configuration Refinement: The ShardConfig class has been refactored to explicitly support FSDP (Fully Sharded Data Parallel) and TP (Tensor Parallel) modes, replacing a generic use_sharding flag with use_fsdp and use_tp boolean parameters. This allows for more granular control over sharding strategies.
• Migration to Native Flax NNX Components: Custom implementations of Einsum and RMSNorm modules have been replaced with their native nnx.Einsum and nnx.RMSNorm counterparts. This streamlines the model architecture by leveraging standard Flax components and their built-in sharding capabilities via kernel_metadata and scale_metadata.
• Updated Model Initialization and Parameter Handling: Model configuration methods (e.g., qwen3_0_6b) now accept use_fsdp and use_tp flags. Parameter mapping in params.py has been adjusted to reflect the new kernel attribute used by nnx.Einsum and nnx.Linear for weights, ensuring compatibility with the updated module structure.
• Enhanced Cache Management and Sharding Integration: The LayerCache now directly utilizes out_sharding for nnx.Cache and nnx.Variable initialization, and updates to cache values are performed using .set_value() for proper nnx state management. The shard utility function has been updated to handle optional PartitionSpec.
• Introduction of Sharding Tests: A new test file, test_sharding_qwen3.py, has been added to specifically verify the sharding functionality of the Qwen3 model, ensuring correct behavior in distributed environments. The test runner script has been updated to execute these tests with appropriate JAX environment settings.
• Flax Dependency Update: The pyproject.toml file has been modified to fetch the Flax dependency directly from its GitHub repository, indicating a reliance on the latest development version of Flax for these changes.

🧠 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/qwen3/modeling.py
  • Replaced functools.partial import with enum.Enum.
  • Introduced ShardMode enum for FSDP and TP.
  • Updated ShardConfig attributes to PartitionSpec | None and modified no_sharding and default methods to accept use_fsdp and use_tp.
  • Modified ModelConfig._from_param and model size methods to pass use_fsdp and use_tp.
  • Adjusted the shard function to handle optional PartitionSpec and s is not None check.
  • Removed custom Einsum and RMSNorm classes.
  • Updated LayerCache to use out_sharding for nnx.Cache and nnx.Variable initialization and .set_value() for updates.
  • Refactored Attention and MLP to use nnx.Einsum and nnx.Linear with kernel_metadata for sharding, and nnx.RMSNorm.
  • Changed property access for head_dim, num_heads, num_kv_heads in Attention to use .kernel.shape.
  • Modified DecoderLayer to initialize nnx.RMSNorm with scale_metadata.
  • Updated Qwen3 to use nnx.Embed with embedding_metadata, nnx.RMSNorm, and nnx.Einsum with kernel_metadata.
  • Moved init_cache from a method of Qwen3 to a standalone function.
  • Adjusted Qwen3.__call__ and forward function signatures and internal calls.
• bonsai/models/qwen3/params.py
  • Updated parameter mapping keys from .w to .kernel for attention projection layers (q_proj, k_proj, v_proj, o_proj) and lm_head to align with nnx.Einsum and nnx.Linear.
• bonsai/models/qwen3/tests/run_model.py
  • Corrected a typo in a comment from 'mtuliple' to 'multiple'.
  • Updated the call to init_cache to use the new standalone modeling.init_cache function.
• bonsai/models/qwen3/tests/test_outputs_qwen3.py
  • Removed the use_sharding=False argument from modeling.ModelConfig.qwen3_0_6b() call.
  • Updated the _init_nnx_cache method to call the standalone modeling.init_cache function.
• bonsai/models/qwen3/tests/test_sharding_qwen3.py
  • Added a new test file to verify sharding functionality.
  • Implemented setUpClass to configure JAX mesh and initialize a sharded Qwen3 model.
  • Included test_full method to check sharding of model outputs.
• pyproject.toml
  • Modified the Flax dependency to point to the Google Flax GitHub repository.
• scripts/run_selective_tests.py
  • Added logic to detect and run sharding-specific tests with JAX_NUM_CPU_DEVICES=4 and JAX_PLATFORMS=cpu environment variables.

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 refactors the Qwen3 model to leverage nnx.Einsum and nnx.RMSNorm, removing custom implementations and simplifying the codebase. The sharding configuration has been made more explicit by replacing use_sharding with use_fsdp and use_tp. A new sharding test has also been added.

My review has identified a couple of critical issues:

• Several model configuration factory methods (qwen3_4b, qwen3_8b, qwen3_14b) are incorrectly calling the ModelConfig constructor directly, which will lead to a TypeError and prevent sharding from being configured.
• There is a duplicated block of code for self.q_proj initialization in the Attention module.

Please address these issues to ensure the model functions correctly.

[vfdev-5]

Put init_cache as module method similar to Gemma3 model