mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
Fix garbage generation preview in hf_ptq.py when pad_token == eos_token (#1673)
### What does this PR do? Type of change: Bug fix <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> Fixes the generation **preview** in `examples/llm_ptq/hf_ptq.py` producing garbage output (e.g. repeated `\u200b` zero-width-space tokens) for models whose tokenizer has `pad_token == eos_token` — most visibly GLM-5.1. The garbage appeared *before* quantization, so it was not a quantization issue. **Root cause:** `pre_quantize` / `post_quantize` take the first (left-padded) calibration sample and call `full_model.generate(preview_input_ids, ...)` **without an `attention_mask`**. HuggingFace only auto-infers the mask when `pad_token_id != eos_token_id` (`generation/utils.py:_prepare_attention_mask_for_generation`); when they are equal it falls back to an all-ones mask, so the model attends to the leading pad/eos tokens, ignores the real prompt, and (for GLM's MoE/DSA/MTP path) collapses to a single repeated token. Calibration itself was always correct — it already passes the mask; only the preview generation was missing it. **Fix:** thread the calibration batch's `attention_mask` through to both preview `generate()` calls. One file changed (`examples/llm_ptq/hf_ptq.py`, +20/-8). ### Usage No usage change — the same command now produces a coherent preview instead of `\u200b` repetition ### Testing Reproduced the exact mechanism (left padding + pad_token == eos_token + missing attention_mask) on a small model(GPT2): without the mask the model emits the same HF warning as the bug report and ignores the prompt; with the mask the output is byte-identical to the unpadded baseline. Verified no behavioral change for models where pad != eos (the explicit mask equals HF's inferred input_ids.ne(pad_id)) and for Whisper (its batch carries no attention_mask, so the path is unchanged). Pre-commit: ruff-check, ruff-format, and mypy (no new errors vs. main) all pass. Before your PR is "Ready for review" Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S). Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.). - Is this change backward compatible?: ✅ <!-- Only changes internal helper signatures within the example script; no public API affected. --> - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A <!-- No copied code, no new dependency. --> - Did you write any new necessary tests?: N/A <!-- Preview path requires model loading; no existing unit-test harness covers it. Verified via a standalone repro of the root-cause mechanism. --> - Did you update Changelog?: N/A <!-- Bug fix confined to an example-script preview; not a library/API change. Happy to add a 0.46 bug-fix entry if preferred. --> - Did you get Claude approval on this PR?: ✅ <!-- Will run `/claude review` before requesting review. --> ### Additional Information Backward compatible across model familes: | Model class | Before (no mask passed) | After (mask passed) | Result | |---|---|---|---| | `pad != eos` (most: T5, BART, many LLMs) | HF infers mask = `input_ids.ne(pad_id)` | explicit calib mask = same tensor | **Identical output** — no change | | `pad == eos` (GLM-5.1, GPT-2-style) | all-ones fallback → attends to pad → garbage | correct mask | **Fixed** | | Whisper | no mask | batch has no `attention_mask` key → `None` → no mask | **Identical** — no change | | Nemotron-VL / DeepSeek / NemotronH / `--skip_generate` | `generate()` not called on this path | unchanged | No change | <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Release Notes * **Bug Fixes** * Enhanced LLM post-quantization example to properly handle attention masks during preview generation. The quantization preview now correctly threads attention masks through generate() calls, ensuring accurate generation outputs are captured both before and after quantization steps. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
This commit is contained in:
@@ -817,13 +817,15 @@ def pre_quantize(
|
||||
"""
|
||||
# Offline specdec models skip pre-quantize preview (no tokenizer or standard dataloader)
|
||||
if args.specdec_offline_dataset is not None:
|
||||
return None, None
|
||||
return None, None, None
|
||||
|
||||
# Only run single sample for preview
|
||||
assert calib_dataloader is not None, "calib_dataloader is required for pre-quantize preview"
|
||||
preview_input_ids = next(iter(calib_dataloader))[
|
||||
"input_features" if model_type == "whisper" else "input_ids"
|
||||
][0:1]
|
||||
batch = next(iter(calib_dataloader))
|
||||
input_key = "input_features" if model_type == "whisper" else "input_ids"
|
||||
preview_input_ids = batch[input_key][0:1]
|
||||
# Pass attention_mask to generate(): HF cannot infer it when pad_token == eos_token.
|
||||
preview_attention_mask = batch["attention_mask"][0:1] if "attention_mask" in batch else None
|
||||
|
||||
# Generate preview before quantization
|
||||
if args.skip_generate:
|
||||
@@ -847,9 +849,13 @@ def pre_quantize(
|
||||
trust_remote_code=args.trust_remote_code,
|
||||
)
|
||||
else:
|
||||
generated_ids_before_ptq = full_model.generate(preview_input_ids, max_new_tokens=100)
|
||||
generated_ids_before_ptq = full_model.generate(
|
||||
preview_input_ids,
|
||||
attention_mask=preview_attention_mask,
|
||||
max_new_tokens=100,
|
||||
)
|
||||
|
||||
return preview_input_ids, generated_ids_before_ptq
|
||||
return preview_input_ids, preview_attention_mask, generated_ids_before_ptq
|
||||
|
||||
|
||||
def post_quantize(
|
||||
@@ -860,6 +866,7 @@ def post_quantize(
|
||||
tokenizer: PreTrainedTokenizerBase | None,
|
||||
processor: ProcessorMixin | None,
|
||||
preview_input_ids,
|
||||
preview_attention_mask,
|
||||
generated_ids_before_ptq,
|
||||
is_nemotron_vl_model,
|
||||
first_text_speech_dataset,
|
||||
@@ -904,7 +911,11 @@ def post_quantize(
|
||||
pass
|
||||
elif model_type != "llama4" and not is_nemotron_vl_model:
|
||||
# Our fake quantizer may not be fully compatible with torch.compile.
|
||||
generated_ids_after_ptq = full_model.generate(preview_input_ids, max_new_tokens=100)
|
||||
generated_ids_after_ptq = full_model.generate(
|
||||
preview_input_ids,
|
||||
attention_mask=preview_attention_mask,
|
||||
max_new_tokens=100,
|
||||
)
|
||||
elif is_nemotron_vl_model and tokenizer is not None:
|
||||
generated_ids_after_ptq = run_nemotron_vl_preview(
|
||||
full_model,
|
||||
@@ -1062,7 +1073,7 @@ def quantize_main(
|
||||
# Detect if this is a Nemotron VL model using architecture-based detection
|
||||
is_nemotron_vl_model = is_nemotron_vl(full_model)
|
||||
|
||||
preview_input_ids, generated_ids_before_ptq = pre_quantize(
|
||||
preview_input_ids, preview_attention_mask, generated_ids_before_ptq = pre_quantize(
|
||||
args, full_model, model_type, tokenizer, calib_dataloader, is_nemotron_vl_model
|
||||
)
|
||||
|
||||
@@ -1173,6 +1184,7 @@ def quantize_main(
|
||||
tokenizer,
|
||||
processor,
|
||||
preview_input_ids,
|
||||
preview_attention_mask,
|
||||
generated_ids_before_ptq,
|
||||
is_nemotron_vl_model,
|
||||
first_text_speech_dataset,
|
||||
|
||||
Reference in New Issue
Block a user