support MTP module WOQ quantization and step3p5 model quantization#1526
Conversation
Signed-off-by: Xin He <xin3.he@intel.com>
|
@copilot Please help add a simple unittest for this logic. |
There was a problem hiding this comment.
Pull request overview
Adds support for preserving auxiliary checkpoint parameters (not instantiated by transformers model classes, e.g. MTP) when exporting/saving quantized models, so they aren’t silently dropped from the output directory.
Changes:
- Introduces
copy_missing_tensors_from_sourceto detect and copy “missing” tensors (by name-prefix components) from the original checkpoint into the exported quantized folder, with optional FP8 dequant and WOQ RTN re-quantization. - Wires the missing-tensor copy step into the export save path and re-exports the helper from
auto_round.utils. - Adds a duplicate-save guard in shard writing and removes now-unused imports.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| auto_round/utils/model.py | Removes unused imports related to prior FP8 handling. |
| auto_round/utils/missing_tensors.py | New implementation for detecting/copying missing tensors, FP8 dequant, and WOQ RTN packing. |
| auto_round/utils/init.py | Re-exports copy_missing_tensors_from_source and adjusts utils exports. |
| auto_round/export/utils.py | Calls missing-tensor copy after save_pretrained. |
| auto_round/compressors/shard_writer.py | Skips saving duplicate tensor names during sharded saves. |
You can also share your feedback on Copilot code review. Take the survey.
|
is there anyway to control the bits for mtp? |
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: xin3he <83260933+xin3he@users.noreply.github.com>
for more information, see https://pre-commit.ci
…saved set (#1528) Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: xin3he <83260933+xin3he@users.noreply.github.com>
Signed-off-by: Xin He <xin3.he@intel.com>
Signed-off-by: Xin He <xin3.he@intel.com>
I add --layer_config argument for it, please check. Verified with vLLM |
Hello, I got an error running this command: I can confirm the build is Is your |
|
Thank you for your attention. I think it's possible that the Transformers download skipped the MTP parameter download because the model structure doesn't include MTP. I suggest using |
Signed-off-by: Xin He <xin3.he@intel.com>
Signed-off-by: Xin He <xin3.he@intel.com>
Signed-off-by: Xin He <xin3.he@intel.com>
Signed-off-by: Xin He <xin3.he@intel.com>
Signed-off-by: Xin He <xin3.he@intel.com>
Signed-off-by: Xin He <xin3.he@intel.com>
Could we add a checker for this? From the model config, we should be able to determine whether the model supports MTP. If the model state dict does not contain the corresponding tensor, we could log a warning or an error to suggest user to use hf download |
|
Command to quantize mixed int4/8 model for Step3p5: |
Signed-off-by: Xin He <xin3.he@intel.com>
|
@Saren-Arterius It's truly a bug, and it's fixed now. |
Signed-off-by: Xin He <xin3.he@intel.com>
Signed-off-by: Xin He <xin3.he@intel.com>
yiliu30
left a comment
There was a problem hiding this comment.
As other reviewers noted that this PR includes several changes.
It would be great, if we can spilit into:
- Add Step 35 support
- Copy MTP param
- Optimizes shard writer
Signed-off-by: Xin He <xin3.he@intel.com>
|
Verified with |
It's fixed. |
Thank you. will keep currenct change due to complex context. |
Really nice work, thanks! |
Signed-off-by: Xin He <xin3.he@intel.com>
|
@xin3he: I updated to the latest build today that includes this patch (42fdb06) and noticed that while I can successfully convert Gemma 3 models with this patch (e.g., unsloth/gemma-3-4b-it), they can no longer be loaded with vLLM v0.17.1 once converted. Reverting to the prior commit (837888d) causes the issue to go away. Here is the relevant error while loading in vLLM with the successfully quantized model: (EngineCore_DP0 pid=2325627) loaded_params = module_load_weights(weights)
(EngineCore_DP0 pid=2325627) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=2325627) File "/home/daved/miniforge3/envs/vllm/lib/python3.12/site-packages/vllm/model_executor/models/siglip.py", line 950, in load_weights
(EngineCore_DP0 pid=2325627) param = params_dict[name]
(EngineCore_DP0 pid=2325627) ~~~~~~~~~~~^^^^^^
(EngineCore_DP0 pid=2325627) KeyError: 'vision_model.encoder.layers.0.self_attn.qkv_proj.weight'
Loading safetensors checkpoint shards: 17% Completed | 1/6 [00:00<00:02, 1.85it/s]Happy to provide more details! |
Description
Background
When
transformersloads a model, it only instantiates the sub-modulesdeclared in the model class. Some architectures include auxiliary
parameters ?? most notably MTP (Multi-Token Prediction) layers ?? that are
stored in the original checkpoint but are not part of the default module
graph. As a result,
model.save_pretrainedsilently omits those tensorsfrom the quantized output.
What's changed?
This module provides :func:
copy_missing_tensors_from_source, which detectsand copies such tensors from the original checkpoint into the saved output
directory. Detection relies on a user-supplied list of name-component
prefixes (
missing_param_prefix, default["mtp"]): a source tensor isconsidered "missing" when it is absent from the saved output and at least
one dot-separated segment of its name starts with one of those prefixes.
Before writing, the function optionally:
weight_scale_invtensor is present in the source checkpoint), and
the saved model (when
quant_method == "auto-round"andpacking_format == "auto_round:auto_gptq"are detected inconfig.json).Additional support step3p5 since it requires mtp
Type of Change
Related Issues
Fixes or relates to #1477, #1534 #1529
Checklist Before Submitting