fix expert_parallel_size to not pass through to vLLM args

[khatwanimohit]

Fix (tunix/generate/vllm_sampler.py): Added args.pop("expert_parallel_size", None) at the top of _vllm_config, immediately after copying _processed_engine_kwargs. This strips the key before it can
reach LLM(**self.args) / EngineArgs(**engine_kwargs), while ep is still correctly derived via resolve_parallelism_sizes and placed into additional_config["sharding"]["sharding_strategy"]["expert_parallelism"].

Test (tests/generate/vllm_sampler_test.py): Added test_expert_parallel_size_via_engine_kwargs_not_leaked_to_vllm which passes expert_parallel_size=2 via engine_kwargs and asserts it does not appear as a top-level key in sampler.args — confirmed failing before the fix (the key was present in args).

Reference

Colab Notebook

Checklist

• I have added all the necessary unit tests for my change.
• I have verified that my change does not break existing code and all unit tests pass.
• I have added all appropriate doc-strings/documentation.
• My PR is based on the latest changes of the main branch (if unsure, rebase the code).
• I have signed the Contributor License Agreement.
• I have followed Contribution Guidelines.

[gemini-code-assist]

Summary of Changes

Hello, 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 resolves an issue where the Tunix-specific expert_parallel_size configuration was being incorrectly passed directly to the vLLM engine arguments, which does not recognize this parameter. The changes ensure that this parameter is properly consumed and translated by Tunix's internal configuration processing, preventing runtime errors and maintaining compatibility with vLLM's API. A new regression test has been added to validate this behavior.

Highlights

• Configuration Handling: Modified the _vllm_config method to explicitly remove expert_parallel_size from the arguments passed to vLLM's EngineArgs, preventing potential errors as this parameter is not directly supported by vLLM.
• Regression Testing: Introduced a new test test_expert_parallel_size_via_engine_kwargs_not_leaked_to_vllm to ensure that expert_parallel_size is correctly processed by Tunix and not inadvertently forwarded to vLLM.

🧠 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

• tests/generate/vllm_sampler_test.py
  • Added test_expert_parallel_size_via_engine_kwargs_not_leaked_to_vllm to verify that expert_parallel_size is not leaked to vLLM.
• tunix/generate/vllm_sampler.py
  • Removed expert_parallel_size from the args dictionary before passing it to vLLM's EngineArgs.

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 correctly fixes a bug where the expert_parallel_size parameter was being improperly passed to the vLLM engine, which would cause a crash. The fix involves removing the parameter from the arguments dictionary before it's passed to vLLM. A comprehensive regression test has been added to verify the fix and prevent this issue from recurring. The changes are logical and well-implemented. I have one minor suggestion regarding a comment to better align with the repository's style guide.

[wang2yn84]

Can you squash the commits please?