Ddp v3

[HDCharles]

SUMMARY:
"please provide a brief summary"

TEST PLAN:
"please outline how the changes were tested"

[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 @HDCharles, 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 introduces support for Distributed Data Parallel (DDP) in AutoRound quantization workflows. It provides a concrete example script showcasing how to quantize a Qwen3 model using DDP and modifies the core AutoRound modifier to correctly handle device mapping when operating in a distributed setting. These changes enable users to leverage multiple GPUs or nodes for faster and more scalable quantization of large language models.

Highlights

• DDP Example Script Added: A new example script, ddp_qwen3_example.py, has been added to demonstrate how to perform AutoRound quantization with Distributed Data Parallel (DDP) using torchrun.
• DDP Device Mapping for AutoRound: The AutoRoundModifier was updated to include a new method, _update_device_map_for_dp, which dynamically sets the device_map for AutoRound based on the current DDP rank, ensuring correct device assignment in a distributed environment.
• Distributed Model Loading: The example script now utilizes init_dist() and load_offloaded_model() for loading models in a DDP setup, allowing for efficient distributed model initialization and offloading.

🧠 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

• experimental/ddp/ddp_qwen3_example.py
  • Added a new script for DDP-enabled AutoRound quantization of Qwen3 models.
  • Included argument parsing for model, scheme, iterations, samples, and deterministic mode.
  • Implemented distributed model loading using init_dist and load_offloaded_model.
  • Integrated AutoRound modifier and oneshot quantization with DDP considerations.
  • Added sample generation and saving of the quantized model in a DDP-aware manner.
• src/llmcompressor/modifiers/autoround/base.py
  • Implemented _update_device_map_for_dp method to set device_map based on DDP rank.
  • Integrated _update_device_map_for_dp call within the apply_autoround function to ensure proper device assignment during distributed quantization.

Activity

• The pull request description indicates that a brief summary and test plan are requested, but no specific details were provided in the initial description.

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

The pull request introduces DDP (Distributed Data Parallel) functionality to the AutoRound quantization example. This involves adding a new example script ddp_qwen3_example.py and modifying src/llmcompressor/modifiers/autoround/base.py to handle device mapping for distributed processing. The changes seem to correctly integrate DDP with the AutoRound modifier, allowing for distributed quantization. However, there are a few areas for improvement regarding code clarity and consistency, particularly in variable naming and import statements.

[gemini-code-assist]

The variable model_name is used here but not defined. It should likely be model_id to match the args.model assignment.

Suggested change

	tokenizer = AutoTokenizer.from_pretrained(model_name)
	tokenizer = AutoTokenizer.from_pretrained(model_id)

[gemini-code-assist]

The torch.distributed as dist import is duplicated. It's already imported on line 20.

Suggested change

	import torch.distributed as dist
	from llmcompressor import oneshot