Add quantization support for DiffusionGemma - #1935
lvliang-intel wants to merge 15 commits into
Conversation
Signed-off-by: lvliang-intel <liang1.lv@intel.com>
…upport_diffusiongemma
for more information, see https://pre-commit.ci
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
Transformers backend test result:test_diffusiongemma_transformers.py CUDA_VISIBLE_DEVICES=2 python test_diffusiongemma_transformers.py [1/5] Loading config ... [2/5] Loading tokenizer ... [3/5] Loading generation_config ... [4/5] Loading quantized model onto cuda:0 ... [5/5] Running a simple Q&A ... Q: Why is the sky blue? Please answer in one short paragraph. transformers load + inference succeeded |
…/auto-round into lvl/support_diffusiongemma
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
|
|
||
| if modules_to_not_convert and not user_specified: | ||
| _DEFAULT_SKIP_KEYWORDS = ("embed", "embed_tokens", "lm_head", "output_embed", "norm") | ||
| modules_to_not_convert = [ |
There was a problem hiding this comment.
Any large change in this file is risky.
Could you provide more details? From my understanding, everything listed above except lm_head should no longer be a linear layer and therefore should already be handled correctly. As for lm_head, it has its own dedicated logic, so I'm not sure why changes would be needed there.
There was a problem hiding this comment.
The original logic is correct for normal models, the get_keys_to_not_convert returns only "lm_head/embed", which aren't quantizable linear layers, so a filter would be a no-op. But DiffusionGemma breaks that assumption because of one regex line in transformers "modeling_diffusion_gemma.py" _tied_weights_keys.
r"encoder.language_model.layers.(?:[^.]+.)*weight" → r"decoder.layers.(?:[^.]+.)*weight"
This ties every .weight under the encoder layers to the matching decoder weight, i.e. self_attn.q_proj.weight, k_proj.weight, v_proj.weight, o_proj.weight, and the per-expert MLP weights. These are exactly the real nn.Linear layers we want to quantize. So without the filter, essentially every encoder and decoder attention/MLP linear is excluded from quantization, it's a bug. The filter restricts the auto-derived skip list back to "embed/lm_head/norm", which is what get_keys_to_not_convert was meant to return.
Do you have better idea (need to be simple) to fix this issue?
|
could you leverage auto-round-best tune a model and upload to intel space, thanks! |
yes, WIP. |
Signed-off-by: lvliang-intel <liang1.lv@intel.com>
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
…upport_diffusiongemma
…/auto-round into lvl/support_diffusiongemma
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
…upport_diffusiongemma Signed-off-by: lvliang-intel <liang1.lv@intel.com>
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
| name for name in modules_to_not_convert if any(key in name for key in _DEFAULT_SKIP_KEYWORDS) | ||
| ] | ||
|
|
||
| if modules_to_not_convert: |
There was a problem hiding this comment.
module_to_not_convert is mainly used for the AWQ format. In vLLM and SGLang, we only parse extra_config.
Besides, if this is caused by the tied-weight key issue, how about handling tied-weight keys explicitly instead?
|
|
||
| if modules_to_not_convert: | ||
| # Pre-compile patterns once instead of recompiling them for every layer name. | ||
| compiled_patterns = [re.compile(n) for n in modules_to_not_convert] |
There was a problem hiding this comment.
@WeiweiZhang1 @n1ck-guo please review this change. I have some concerns that it might be risky.
There was a problem hiding this comment.
From the code, the risk is not significant. The only issue is that the impact is large. Could we limit it to a controllable range using methods like patching, or perhaps control it through parameters?
There was a problem hiding this comment.
If the user mark some layers as non-quantized, there should be no entries in module_not_convert when the format is AutoRound instead of AWQ. In AutoRound, we only use extra_config to indicate whether the lm-head is quantized or to mark certain layers as not quantized.
So I think handling this logic via module_not_convert might be problematic.
There was a problem hiding this comment.
Agree with this point. I will change the code as below.
# AWQ format: exclude specified modules.
if "awq" in (getattr(quantization_config, "quant_method", "") or "").lower():
extra_config = skip_not_convert_modules(model, quantization_config, layer_names, extra_config)
Signed-off-by: lvliang-intel <liang1.lv@intel.com>
Signed-off-by: lvliang-intel <liang1.lv@intel.com>
…/auto-round into lvl/support_diffusiongemma
|
This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you! |
4 participants
Description
Add quantization support for DiffusionGemma (transformers ≥ 5.11.0). The model declares encoder/decoder MoE gate_up_proj / down_proj as tied across stacks; AutoRound's MoE-unfuse step splits these fused 3D parameters into per-expert linears, which leaves the declared _tied_weights_keys referring to nonexistent parameter names and breaks downstream loading/export.
Type of Change
New feature
Related Issues
None
Checklist Before Submitting
/azp run Unit-Test-CUDA-AutoRound.