[DFX] Adaptable to multiple model validations for fused moe

[kaniel-outis]

Adaptable to multiple model validations for fused moe:

1. adapt longcat （hidden 6144，hidden-mid 2048）
2. adapt Qwen3-235B-A22B（hidden 4096，hidden-mid 1536）

[gemini-code-assist]

Summary of Changes

Hello @kaniel-outis, 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 enhances the testing framework for fused Mixture-of-Experts (MoE) models by introducing greater adaptability for different model architectures. The core change involves making the intermediate dimension (hidden_mid) a configurable parameter throughout the test suite, from argument parsing to weight initialization and test execution. This allows for more flexible validation against diverse model specifications, such as the 'longcat' configuration mentioned, ensuring the MoE implementation can handle varying hidden and intermediate layer sizes.

Highlights

• Configurable Intermediate Dimension: Introduced a new hidden_mid parameter across test functions (test, test_loop) and argument parsing, allowing the intermediate dimension of the fused MoE model to be specified and varied during testing.
• Dynamic Output Dimension Calculation: Modified the init_base_weights function to dynamically calculate hidden_out as hidden_mid // 2, removing it as a direct input parameter and making the output dimension dependent on the intermediate dimension.
• Flexible Weight Reshaping: Updated the reshape_fusion_gmm_weight function to use -1 for an inferred dimension during weight reshaping, enhancing flexibility to accommodate different hidden_mid values without hardcoding.
• Command-Line Argument for hidden_mid: Added a --hidden-mid command-line argument to the test script, enabling users to easily configure the intermediate dimension when running tests for various model configurations.

Ignored Files

• Ignored by pattern: .github/workflows/** (1)
  • .github/workflows/pr-test-npu.yml

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 adapts the fused MoE tests to support multiple model validations, specifically for different hidden dimensions. This is achieved by making hidden_mid a configurable parameter, calculating hidden_out based on hidden_mid, and generalizing a weight reshaping function. The changes are logical and consistent. I've provided a couple of suggestions to improve documentation accuracy and code clarity.