Conversation
Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com> Co-authored-by: yiliu30 <yi4.liu@intel.com>
Signed-off-by: yiliu30 <yi4.liu@intel.com>
Signed-off-by: yiliu30 <yi4.liu@intel.com>
|
👋 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. |
Summary of ChangesHello, 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 enhances the Highlights
Changelog
Using Gemini Code AssistThe 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
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for Distributed Data Parallel (DDP) with AutoRound, demonstrated by the new example script ddp_qwen3_example.py. A security audit identified a low-risk path traversal concern in an example script due to user-supplied arguments for SAVE_DIR, though it is not considered a critical issue in this context. The AutoRoundModifier correctly handles device mapping in a distributed environment. One suggestion has been provided to improve the structure of the new example script to align with Python best practices and prevent potential import errors.
| if __name__ == "__main__": | ||
| parser = argparse.ArgumentParser( | ||
| description="AutoRound Quantization with DDP support" | ||
| ) | ||
| parser.add_argument( | ||
| "--model", | ||
| type=str, | ||
| default="Qwen/Qwen3-8B", | ||
| help="Model name or path", | ||
| ) | ||
| parser.add_argument( | ||
| "--scheme", | ||
| type=str, | ||
| default="W4A16", | ||
| help="Quantization scheme (W4A16, MXFP8, MXFP4, etc.)", | ||
| ) | ||
| parser.add_argument("--iters", type=int, default=200, help="Number of iterations") | ||
| parser.add_argument("--nsamples", type=int, default=128, help="Number of samples") | ||
| parser.add_argument( | ||
| "--disable_torch_compile", | ||
| action="store_true", | ||
| help="Disable torch.compile for model acceleration during quantization", | ||
| ) | ||
| parser.add_argument( | ||
| "--deterministic", | ||
| action="store_true", | ||
| help="Enable deterministic mode for reproducibility", | ||
| ) | ||
| args = parser.parse_args() | ||
|
|
||
| if args.deterministic: | ||
| config_deterministic() | ||
|
|
||
| model_id = args.model | ||
|
|
||
| ###### DDP MODEL LOAD CHANGE ##### | ||
| init_dist() | ||
| with load_offloaded_model(): | ||
| model = AutoModelForCausalLM.from_pretrained( | ||
| model_id, dtype="auto", device_map="auto_offload" | ||
| ) | ||
| ################################## | ||
|
|
||
| tokenizer = AutoTokenizer.from_pretrained(model_id) | ||
|
|
||
| # Select calibration dataset. | ||
| NUM_CALIBRATION_SAMPLES = args.nsamples | ||
| MAX_SEQUENCE_LENGTH = 2048 | ||
| ITERS = args.iters | ||
|
|
||
|
|
||
| # Get aligned calibration dataset. | ||
| from auto_round.calib_dataset import get_dataset # noqa: E402 | ||
|
|
||
| # Note: Make sure model are loaded before importing auto-round related code. | ||
| # This requirement will be lifted once switching to new release of auto-round which | ||
| # includes below fix: | ||
| from llmcompressor.modifiers.autoround import AutoRoundModifier # noqa: E402 | ||
|
|
||
| ds = get_dataset( | ||
| tokenizer=tokenizer, | ||
| seqlen=MAX_SEQUENCE_LENGTH, | ||
| nsamples=NUM_CALIBRATION_SAMPLES, | ||
| ) | ||
|
|
||
| # Configure the quantization algorithm to run. | ||
| # * quantize the weights to 4 bit with AutoRound with a group size 128 | ||
| recipe = AutoRoundModifier( | ||
| targets="Linear", | ||
| scheme=args.scheme, | ||
| ignore=[ | ||
| "lm_head", | ||
| "re:.*mlp.gate$", | ||
| ], | ||
| iters=ITERS, | ||
| enable_torch_compile=not args.disable_torch_compile, | ||
| ) | ||
|
|
||
| # Apply algorithms. | ||
| oneshot( | ||
| model=model, | ||
| dataset=ds, | ||
| recipe=recipe, | ||
| max_seq_length=MAX_SEQUENCE_LENGTH, | ||
| num_calibration_samples=NUM_CALIBRATION_SAMPLES, | ||
| shuffle_calibration_samples=False, | ||
| ) | ||
|
|
||
| rank = dist.get_rank() | ||
| logger.info(f"[Rank {rank}] Quantization completed") | ||
| # Confirm generations of the quantized model look sane. | ||
| logger.info("\n\n") | ||
| logger.info("========== SAMPLE GENERATION ==============") | ||
| dispatch_model(model) | ||
| sample = tokenizer("Hello my name is", return_tensors="pt") | ||
| sample = {key: value.to(model.device) for key, value in sample.items()} | ||
| output = model.generate(**sample, max_new_tokens=100) | ||
| logger.info(tokenizer.decode(output[0])) | ||
| logger.info("==========================================\n\n") | ||
|
|
||
| logger.info("Saving...") | ||
| # Save to disk compressed. | ||
| SAVE_DIR = ( | ||
| model_id.rstrip("/").split("/")[-1] | ||
| + f"-{args.scheme}-AutoRound" | ||
| + f"-iters{args.iters}-nsamples{args.nsamples}" | ||
| + "-DDP" | ||
| + str(dist.get_world_size()) | ||
| ) | ||
| model.save_pretrained(SAVE_DIR, save_compressed=True) | ||
| tokenizer.save_pretrained(SAVE_DIR) | ||
| logger.info(f"Saved to {SAVE_DIR}") | ||
|
|
||
| dist.destroy_process_group() |
There was a problem hiding this comment.
To ensure the script is only executed when run directly and to prevent issues with variable scope if imported, it's best practice to place all script logic inside the if __name__ == "__main__": block. This also fixes a bug where args would be undefined if this file were imported elsewhere.
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="AutoRound Quantization with DDP support"
)
parser.add_argument(
"--model",
type=str,
default="Qwen/Qwen3-8B",
help="Model name or path",
)
parser.add_argument(
"--scheme",
type=str,
default="W4A16",
help="Quantization scheme (W4A16, MXFP8, MXFP4, etc.)",
)
parser.add_argument("--iters", type=int, default=200, help="Number of iterations")
parser.add_argument("--nsamples", type=int, default=128, help="Number of samples")
parser.add_argument(
"--disable_torch_compile",
action="store_true",
help="Disable torch.compile for model acceleration during quantization",
)
parser.add_argument(
"--deterministic",
action="store_true",
help="Enable deterministic mode for reproducibility",
)
args = parser.parse_args()
if args.deterministic:
config_deterministic()
model_id = args.model
###### DDP MODEL LOAD CHANGE #####
init_dist()
with load_offloaded_model():
model = AutoModelForCausalLM.from_pretrained(
model_id, dtype="auto", device_map="auto_offload"
)
##################################
tokenizer = AutoTokenizer.from_pretrained(model_id)
# Select calibration dataset.
NUM_CALIBRATION_SAMPLES = args.nsamples
MAX_SEQUENCE_LENGTH = 2048
ITERS = args.iters
# Get aligned calibration dataset.
from auto_round.calib_dataset import get_dataset # noqa: E402
# Note: Make sure model are loaded before importing auto-round related code.
# This requirement will be lifted once switching to new release of auto-round which
# includes below fix:
from llmcompressor.modifiers.autoround import AutoRoundModifier # noqa: E402
ds = get_dataset(
tokenizer=tokenizer,
seqlen=MAX_SEQUENCE_LENGTH,
nsamples=NUM_CALIBRATION_SAMPLES,
)
# Configure the quantization algorithm to run.
# * quantize the weights to 4 bit with AutoRound with a group size 128
recipe = AutoRoundModifier(
targets="Linear",
scheme=args.scheme,
ignore=[
"lm_head",
"re:.*mlp.gate$",
],
iters=ITERS,
enable_torch_compile=not args.disable_torch_compile,
)
# Apply algorithms.
oneshot(
model=model,
dataset=ds,
recipe=recipe,
max_seq_length=MAX_SEQUENCE_LENGTH,
num_calibration_samples=NUM_CALIBRATION_SAMPLES,
shuffle_calibration_samples=False,
)
rank = dist.get_rank()
logger.info(f"[Rank {rank}] Quantization completed")
# Confirm generations of the quantized model look sane.
logger.info("\n\n")
logger.info("========== SAMPLE GENERATION ==============")
dispatch_model(model)
sample = tokenizer("Hello my name is", return_tensors="pt")
sample = {key: value.to(model.device) for key, value in sample.items()}
output = model.generate(**sample, max_new_tokens=100)
logger.info(tokenizer.decode(output[0]))
logger.info("==========================================\n\n")
logger.info("Saving...")
# Save to disk compressed.
SAVE_DIR = (
model_id.rstrip("/").split("/")[-1]
+ f"-{args.scheme}-AutoRound"
+ f"-iters{args.iters}-nsamples{args.nsamples}"
+ "-DDP"
+ str(dist.get_world_size())
)
model.save_pretrained(SAVE_DIR, save_compressed=True)
tokenizer.save_pretrained(SAVE_DIR)
logger.info(f"Saved to {SAVE_DIR}")
dist.destroy_process_group()2 participants
SUMMARY:
"please provide a brief summary"
TEST PLAN:
"please outline how the changes were tested"