add profile for qwen_model #114
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @aolemila, I'm Gemini Code Assist1! 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 primarily focuses on integrating JAX profiling capabilities into the QWen model's test suite. The goal is to enable detailed performance analysis of the QWen model's decoding process. Additionally, it includes a minor correction to an expected output shape in a test and resolves a previous merge conflict.
Highlights
- JAX Profiling Integration: I've added JAX profiler tracing to the
test_qwen_model_decodemethod intest_qwen_model.py. This will generate performance profiles for the QWen model's decoding process, saving them to/root/users/aolemila/jax_profile_sglang_qwen/profile. - Test Adjustment for Profiling Accuracy: To ensure accurate profiling, I've added
x.block_until_ready()within thetest_qwen_model_decodemethod. This explicitly waits for JAX computations to complete before proceeding, ensuring all operations are captured by the profiler. - Vocabulary Size Correction: The expected vocabulary size in
test_qwen_model_prefillwas updated from 10000 to 10048. This corrects an assertion to match the actual output shape of the QWen model's logits. - Merge Conflict Resolution: A merge conflict related to the import of
QWenLMHeadModelvs.QWenLMHeadJaxModelintest_qwen_model.pyhas been resolved, standardizing onQWenLMHeadJaxModel.
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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| 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 issue 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 is currently in preview and 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 to provide feedback.
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
-
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. ↩
There was a problem hiding this comment.
Code Review
This pull request adds a profile for the Qwen model, including changes to the vocabulary size in the test, adding a JAX profiler trace, and adding a block until ready call. The changes seem reasonable, but the vocabulary size change and the profiler trace should be verified.
| @@ -82,10 +78,10 @@ def test_qwen_model_prefill(self): | |||
| forward_batch.positions, forward_batch) | |||
| # Now y is LogitsProcessorOutput with next_token_logits for each sequence | |||
| # Shape: [batch_size, vocab_size] = [128, 10000] | |||
| self.assertEqual(y.next_token_logits.shape, (128, 10000)) | |||
| self.assertEqual(y.next_token_logits.shape, (128, 10048)) | |||
There was a problem hiding this comment.
The vocabulary size has changed. Ensure this change is intentional and that the model is correctly configured to handle the new vocabulary size.
self.assertEqual(y.next_token_logits.shape, (128, 10048))|
|
||
| def test_qwen_model_decode(self): | ||
| with self.mesh: | ||
| with self.mesh, jax.profiler.trace("/root/users/aolemila/jax_profile_sglang_qwen/profile", create_perfetto_link=True): |
There was a problem hiding this comment.
Adding a JAX profiler trace is useful for debugging and performance analysis. However, ensure that this profiling code is removed or disabled before merging into a production environment to avoid performance overhead.
with self.mesh, jax.profiler.trace("/root/users/aolemila/jax_profile_sglang_qwen/profile", create_perfetto_link=True):| @@ -123,6 +119,7 @@ def test_qwen_model_decode(self): | |||
| decoded_token = tokenizer.decode([current_token_id]) | |||
| print( | |||
| f"Step {i+1}: token_id={current_token_id}, decoded='{decoded_token}'") | |||
| x.block_until_ready() | |||
There was a problem hiding this comment.
Adding x.block_until_ready() will ensure that the computation is complete before proceeding. This is useful for debugging and benchmarking, but can hurt performance if used unnecessarily. Consider if this block is needed in the long run.
x.block_until_ready()
Motivation
Modifications
Checklist