[Bugfix][Gemma4] Fix AutoRound GPTQ router and INC tail shards#41588
[Bugfix][Gemma4] Fix AutoRound GPTQ router and INC tail shards#41588lesj0610 wants to merge 1 commit into
Conversation
Add a correctness fallback for INC/GPTQ row-parallel shards whose TP partition crosses group boundaries, and dequantize Gemma4 AutoRound router projection weights into GateLinear's dense parameter during load.\n\nKeep Gemma4 MoE activation and packed expert name mapping out of this branch because those are covered by separate open PRs.\n\nCo-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: lesj0610 <lesj0610@users.noreply.github.com>
There was a problem hiding this comment.
Code Review
This pull request introduces the INCGPTQRowParallelTailLinearMethod fallback to handle row-parallel GPTQ layers where tensor parallel shards are not aligned with quantization group boundaries, and implements a dequantization-on-load strategy for MoE routers in the Gemma4 model. Feedback was provided to restrict the new fallback method to layers sharded along the input dimension to ensure correct group index calculation. Additionally, it was noted that the router projection should remain unquantized in memory to avoid a KeyError during the weight loading process.
| sym, | ||
| ) | ||
| if ( | ||
| isinstance(layer, LinearBase) |
There was a problem hiding this comment.
The INCGPTQRowParallelTailLinearMethod fallback is specifically designed for row-parallel layers where tensor parallelism can split a quantization group across ranks. Its create_weights implementation (lines 728-731) calculates shard_offset by multiplying the rank by the partition width, which is only correct for RowParallelLinear. For ColumnParallelLinear (and its variants like QKVParallelLinear), the input dimension is not sharded, so shard_offset should be 0 for all ranks. Using this fallback for non-row-parallel layers will result in incorrect group indices (g_idx) and corrupted dequantized weights on ranks > 0. You should restrict this fallback to layers that are actually sharded along the input dimension.
| isinstance(layer, LinearBase) | |
| isinstance(layer, LinearBase) and getattr(layer, "input_is_parallel", False) |
| config.num_experts, | ||
| bias=False, | ||
| out_dtype=torch.float32, | ||
| quant_config=quant_config, |
There was a problem hiding this comment.
The router projection should remain unquantized in memory to support the dequantization-on-load strategy implemented in load_weights. Passing quant_config here causes GateLinear to be instantiated as a quantized layer (e.g., using GPTQLinearMethod), which registers qweight instead of weight. Consequently, the load_weights logic at line 1481 will fail with a KeyError when it attempts to access router_name.weight. Since the router is small, dequantizing it into a dense weight during loading is a reasonable trade-off for performance, but it requires the layer itself to be unquantized.
| quant_config=quant_config, | |
| quant_config=None, |
|
Closing per author workflow correction. This split will be handled in fork-only PRs before any upstream submission. |
Summary
Fix the remaining Gemma4 AutoRound GPTQ runtime pieces that are independent from MoE activation support:
qweight/qzeros/scalesinto the denseGateLinearrouter weight during loadingThis PR intentionally does not include Gemma4 MoE
gelu_tanhactivation wiring or broad packed expert weight-name remapping.Duplicate-work check
Checked open PRs before opening:
gelu_tanhactivation wiring. This PR excludes that activation change.Tests
Result: targeted tests passed.
AI assistance was used to prepare and split this change; I reviewed the resulting diff and validation output.