[GPTQ] -- parallelized multi-GPU quantization in compress_modules

[aashvgit]

SUMMARY:
I've added ThreadPoolExecuter To manage parallel threads and also added New Helper Method (quantize_and_update_module and updated the compress_modules making it simple and making ThreadPoolExecuter with one worker for each GPU (OPTION B)

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[gemini-code-assist]

Summary of Changes

Hello @aashvgit, 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 improves the efficiency of the GPTQ quantization process by enabling parallel execution. It refactors the core compression logic to utilize a ThreadPoolExecutor, allowing multiple modules to be quantized concurrently across available GPUs. This enhancement aims to reduce the overall time required for model quantization, with a robust fallback to sequential CPU processing if no GPU resources are detected.

Highlights

• Parallel Quantization: Implemented parallel execution for GPTQ quantization using ThreadPoolExecutor to distribute tasks across available GPUs, significantly speeding up the process.
• New Helper Method: Introduced _quantize_and_update_module to encapsulate the logic for quantizing and updating a single module, facilitating its use in parallel threads.
• GPU Detection & Fallback: Added logic to detect available CUDA devices and utilize them for parallel quantization; if no GPUs are found, the process gracefully defaults to sequential execution on the CPU.
• Refactored compress_modules: The compress_modules method was simplified and updated to orchestrate the new parallel quantization workflow, making it cleaner and more efficient.

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 introduces parallel execution for GPTQ quantization across multiple GPUs using a ThreadPoolExecutor. The implementation refactors the quantization logic into a helper function, which is then called in parallel for different modules. This is a good approach to speed up the quantization process. However, I've found a potential issue with how the original device of a module is handled, which could lead to incorrect behavior in model-parallel scenarios. My feedback includes a specific code suggestion to address this.

[gemini-code-assist]

The current implementation assumes that all modules to be quantized reside on the same device initially. It determines the original_device from the first module in the list and uses it for all modules. This can cause problems in a model-parallel setup where modules might be distributed across different devices. If that's the case, after quantization, all modules would be incorrectly moved to the device of the first module.

To make this more robust, you should determine and store the original device for each module individually and use that to move the module back after processing.

Suggested change

	original_device = get_execution_device(modules_to_quantize[0])
	with ThreadPoolExecutor(max_workers=len(gpu_ids)) as executor:
	futures = []
	for i, module in enumerate(modules_to_quantize):
	target_gpu_id = gpu_ids[i % len(gpu_ids)]
	target_device = f"cuda:{target_gpu_id}"
	future = executor.submit(
	self._quantize_and_update_module,
	module,
	target_device,
	original_device
	)
	comp_logger.set_loss(loss)
	futures.append(future)
	for future in futures:
	future.result()
	original_devices = {m: get_execution_device(m) for m in modules_to_quantize}
	with ThreadPoolExecutor(max_workers=len(gpu_ids)) as executor:
	futures = []
	for i, module in enumerate(modules_to_quantize):
	target_gpu_id = gpu_ids[i % len(gpu_ids)]
	target_device = f"cuda:{target_gpu_id}"
	future = executor.submit(
	self._quantize_and_update_module,
	module,
	target_device,
	original_devices[module],
	)
	futures.append(future)
	for future in futures:
	future.result()

[brian-dellabetta]

This might make sense, if we are allowing multi-gpu compression, modules could also be spread out across different GPUs.

[brian-dellabetta]

Thanks for the contribution! The main maintainer for GPTQ is out of office, but I will discuss this with him next week.

[brian-dellabetta]

This might make sense, if we are allowing multi-gpu compression, modules could also be spread out across different GPUs.

[aashvgit]

Hi i would like to know if there's any additional changes or edge cases to work with i was thinking to add multi-GPU / model-parallel test coverage if that would be helpful. should i proceed like that? @brian-dellabetta

[brian-dellabetta]

Hi @aashvgit , thanks for the contribution and sorry for the slow response. I was on leave the past few weeks. Regarding your question, I left one additional comment and the gemini comment looks valid as well.

I will bring this up with the team tomorrow. We are discussing parallelization approaches and I think something like this at the level of the modifier could be a quick win for GPTQ (and most likely AWQ). Do you have any benchmarks from running this? Do you see a nice speedup?

[brian-dellabetta]

we might be able to shorten this using a different part of the concurrent.futures API. Can try this map feature instead of needing another for loop to call future.result()?

import concurrent.futures
import time

def task_function_simple(x):
    time.sleep(1) # Simulate work
    return x * x

def main_map():
    inputs = [1, 2, 3, 4, 5]
    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
        # map automatically submits tasks and yields results in order
        results = executor.map(task_function_simple, inputs)
        print("Tasks submitted. Results in order:")
        for result in results:
            print(f"Result: {result}")

if __name__ == "__main__":
    main_map()

[brian-dellabetta]

Hi @aashvgit , we spoke internally this morning about your PR. We are considering some other pathways for parallelized compression that can be applied more generically to any modifier, not specific to GPTQ, namely data-parallelism and possibly other options. It's a big consideration, and in the new year @kylesayrs will be creating an RFC for users to provide feedback. We will be sure to include this modifier-specific parallelism as one of the options, and will link in your PR. Let's leave this open and we can revisit once we have gathered some user feedback.

We appreciate your taking the time to create this PR. I personally prefer the simplicity of your approach, but we'd have to find ways to add it to each modifier, and we need to gather more user feedback to see what path is most suitable, and we'll have to do more comprehensive benchmarking to see how much speedup can be provided (please let us know if you have any data on this). To be continued!

[aashvgit]

Hello @brian-dellabetta Thank you for the reply and also sorry for the late reply here, i think the current implementation does assumes all modules originate from the same device which i now realised that it's not a right way for multi-GPU setups maybe update the logic to track the original device per module and restore each module to respective device after that instead of using a single device. Would like you guide if the approach is correct or not maybe I'll check it and gather some data and will share the same and thanks for considering it.