fix: quantization on multiple gpus#284
Conversation
|
/gemini review |
![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 @dalongbao, 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 addresses a critical bug in the model quantization process, particularly when operating in multi-GPU environments. The core issue was that model parameters were not consistently placed on their correct devices after quantization. The changes ensure that the device map is accurately updated and used to guide the placement of all model parameters, leading to more robust and correct behavior for quantized models across multiple GPUs.
Highlights
- Multi-GPU Quantization Fix: Corrects a bug where model parameters were not consistently moved to their intended devices after quantization, specifically addressing issues in multi-GPU environments by ensuring proper device assignment.
- Dynamic Device Map Management: The
quantizefunction now dynamically re-infers and updates thedevice_mapusingaccelerateutilities after initial model preprocessing. This ensures the device map accurately reflects the post-quantization state and guides correct parameter placement. - Precise Parameter Device Assignment: During the quantization process, individual parameters are now explicitly assigned to their correct target devices based on the newly inferred
device_map, preventing parameters from being left on incorrect devices.
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 fixes an issue with quantization on multiple GPUs by re-calculating the device map after quantization and ensuring parameters are moved to the correct devices based on this updated map. The changes look good, but I've identified a potential issue in the device lookup logic that could lead to incorrect device placement. My suggestion addresses this to make the implementation more robust.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Code Review
This pull request aims to fix quantization on multiple GPUs by updating the device map after the model has been preprocessed for quantization. While the overall approach of re-inferring the device map is sound, the implementation introduces a critical runtime error due to an incomplete refactoring of a function's call sites. Additionally, there's a high-severity bug in the device placement logic that could lead to incorrect device assignment for model parameters, and a medium-severity maintainability issue with a hardcoded server address. I've provided detailed comments and suggestions to address these points.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request addresses a bug in model quantization when using multiple GPUs. The core of the fix is within the quantize function, which now correctly re-infers the device_map after the model has been preprocessed for quantization. This ensures that model parameters are placed on the correct devices. The function signature for quantize has been updated to return this new device_map, and the calling functions have been adjusted accordingly. The logic appears sound and directly targets the issue. I have one suggestion to improve code clarity by removing a line of dead code.
future-xy
added
the
ready-for-gpu-testing
|
At first glance, this appears to work in a multi-GPU environment. However, after a closer look, I believe we need to double-check the behavior—currently, all quantized parameters are being loaded onto 
Figure explanation: This shows loading OPT-6.7B on 4 GPUs with FP4 precision. Initially, each GPU holds around 3.3GB of FP16 parameters. Later, |
2 participants
Description
Updated device map and param devices after quantization to ensure parameters get moved correctly.
Motivation
Addresses #282
Type of Change
Checklist