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:
Frida Hou
2026-06-10 20:50:30 -07:00
committed by GitHub
parent 129d22d4ab
commit c88b62beec
+20 -8
View File
@@ -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,