support lm-head/embed requant

[ZX-ModelCloud]

No description provided.

[ZX-ModelCloud]

PR Review
PR: #3050
Reviewed head SHA: 319aee2
Result: findings

Findings:

1. Severity: medium
   File and line: optimize/requant_embed_lm_head.py:85
   Rationale: On a tied-weight model, get_output_embeddings() is the same object as the input embedding, and named_modules() normally resolves both endpoints to the input path. The CLI therefore deduplicates the output target here and installs the requested bits/group-size override only for the input path. ModuleLooper unties the modules later, producing a new lm_head path; that path then receives the hard-coded 8-bit/group-32 fallback instead of the requested output configuration. Output-only and both modes can silently quantize the output endpoint with the wrong settings.
   Action: Untie tied endpoints before resolving target_names, or explicitly carry the requested output override from the pre-untie alias to the new output module name. Add a tied-model test that asserts the effective config on both runtime paths.
2. Severity: high
   File and line: gptqmodel/models/writer.py:291
   Rationale: Single-endpoint requantization cannot save a tied source checkpoint correctly. In output-only mode, matched_shards_by_prefix contains only the output prefix, so the fallback condition requiring the input prefix in that mapping is false and save raises at line 304 because tied checkpoints commonly omit lm_head.weight. In input-only mode, only the input replacement is written, but line 785 persists tie_word_embeddings=false after ModuleLooper cloned an untied output head; the cloned lm_head weight is never serialized, so reload has a missing output weight.
   Action: Resolve the original tied input shard independently of the selected replacement prefixes and serialize the untied output state whenever the saved config becomes untied, including input-only mode. Cover input-only and output-only saves from a source containing only the tied input tensor, then reload both outputs.
3. Severity: medium
   File and line: gptqmodel/nn_modules/qlinear/init.py:1549
   Rationale: The embedding branch leaves int_weight with one row per vocabulary entry, but the packing loop always reads a full pack_factor rows from every ceil-sized qweight row. For vocabularies not divisible by pack_factor, such as GPT-2 with 50,257 tokens at 4 bits, the final iteration indexes beyond int_weight and crashes. The current unit test uses a five-row embedding but stops before packing, while the end-to-end checkpoint happens to have an aligned vocabulary.
   Action: Pad the embedding rows and corresponding g_idx to the packing width or bound and zero-fill the final partial word, and ensure dequantization preserves or slices to the original vocabulary size. Add a pack/save/reload/forward test with a non-divisible vocabulary.

[Qubitium]

Reviewed head 319aee222a5bff92c06bfeb7c3ec9b136e490247.

Result: changes recommended — three correctness findings. This review confirms two issues from the earlier discussion and adds the calibrated embed_only issue.

1. [P1] Single-endpoint requantization does not preserve a tied source checkpoint.
   writer.py:291–304
   When the source has tie_word_embeddings=true and stores only the input embedding weight, output-only saving cannot use the fallback: matched_shards_by_prefix contains only selected replacements, so the required input prefix is absent. It raises ValueError: Could not find checkpoint tensor for embedding module lm_head. Input-only saving has the complementary problem: the looper creates an independent output head and the writer saves tie_word_embeddings=false, but only the input replacement is serialized; the source's omitted output weight remains missing. The resulting checkpoint cannot faithfully reload the runtime model.

   Resolve the tied source shard independently of the replacement selection, and include the retained output weight when untying an input-only checkpoint. Add input-only and output-only save/reload coverage using a source that omits the tied output tensor. The new lifecycle test covers only BOTH.

2. [P2] Embedding packing reads past vocabularies that do not fill a packing word.
   qlinear/init.py:1535–1549
   The new embedding branch keeps one int_weight row per token, while the existing packing loop unconditionally reads pack_factor rows for every ceil-sized output word. At 4 bits/int32, a vocabulary of 17 reaches row 17 and raises IndexError; 50,257 rows fails the same way. The embedding class advertises arbitrary input dimensions and automatic padding, but this path does not pad the source rows.

   Zero-fill or safely bound partial words, and keep unpacking/g_idx consistent with the original vocabulary length. Add a pack/save/reload/lookup test with an unaligned vocabulary, rather than stopping at the floating-point quantization result.

3. [P2] The calibrated embedding path silently ignores embed_only=False.
   stage_layer.py:545–548
   quantize() now forwards QuantizeEmbedConfig to ModuleLooper, but the looper retains only embed_quant_mode. This condition consequently disables decoder subset planning whenever any embedding mode is present, even for a fresh unquantized model requested with QuantizeEmbedConfig(..., embed_only=False). Decoder blocks are replayed without quantization, although the operation ultimately marks the model quantized. The weight-only looper already distinguishes this flag.

   Carry embed_only through the calibrated flow and suppress decoder quantization only for the appropriate embedding-only operation. Test a fresh small model with embed_only=False, asserting that both the requested endpoint and decoder projections are quantized.

Validation: inspected the 19-file diff and surrounding quantization, packing, and serialization code. Executed the actual writer helper functions extracted from this commit with in-memory tensor/I/O doubles: OUTPUT raised the missing-prefix error, INPUT wrote only the input replacement, and BOTH wrote both replacements. Executed the actual NumPy allocation/packing loop: vocabulary 16 passed; 17 and 50,257 raised the out-of-bounds error. These are focused logic reproductions, not full checkpoint integration tests. Full PyTorch/CUDA tests were unavailable because PyTorch is not installed in this environment.

Clarification on the earlier CLI finding: tying weights does not by itself make the input and output modules the same object. Distinct embedding/linear modules can share a weight parameter while retaining different runtime paths. The earlier path-deduplication claim needs an actual module-aliasing model to establish its applicability; I am not counting it as a general tied-weight defect here.

[devin-ai-integration]

Review: PR #3050 — support lm-head/embed requant

Reviewed head 319aee22 (branch zx_support_requant vs main). Ran the PR's CPU tests locally: test_requant_embed_lm_head.py test_embed_quant_api.py test_embedding_save_lifecycle.py test_gptq.py → 44 passed, 4 skipped (GPU-only). test_requant_embed_lm_head_e2e.py skipped (needs /monster/data/model/Qwen2.5-0.5B-Instruct-gptq-4bit + CUDA). Note the PR currently has a merge conflict with main and needs a rebase.

Verdict: request changes — the core GPTQ-on-nn.Embedding path and the pack/dequant layout are consistent (checked PackableQuantLinear.pack embedding branch against create_quant_module in_features=num_embeddings / out_features=embedding_dim and TorchQuantEmbeddings.forward), but there are correctness gaps around tied embeddings and the embed_only flag.

High

1. embed_only is silently ignored on the GPTQ/calibration path (stage_layer.py if embed_quant_mode is not None and not is_embeddings_module: subset_plans = []). Any quantize(..., embed_quant_config=QuantizeEmbedConfig(embed_only=False)) on an fp model now skips every decoder layer and only quantizes the endpoints, yet the saved quantize_config claims a fully quantized GPTQ checkpoint. Before this PR _quantize_with_calibration never received embed_quant_config, so this changes behavior for existing callers. Either honor embed_only (quantize layers when False) or raise when embed_only=False reaches ModuleLooper.

2. Tied-embedding checkpoints cannot be saved correctly in single-endpoint modes (writer.py::_save_embedding_replacement_safetensors, agrees with finding 2 of the earlier automated review):

   • OUTPUT mode: prefixes == {lm_head}; a tied checkpoint has no lm_head.weight, so matched_shards is empty, and tied_output_fallback is None because the input prefix is not in matched_shards_by_prefix → ValueError("Could not find checkpoint tensor for embedding module lm_head"). The fallback condition should look up the input shard from turtle_model._weight_map directly, not from the selected prefixes.
   • INPUT mode: ModuleLooper calls untie_word_embeddings for every embed_quant_mode, so config.tie_word_embeddings is written as false by ModelWriter, but the cloned fp lm_head.weight is never serialized → reload has no output weight. Either only untie when the output endpoint is requested, or also write the untied lm_head state whenever the saved config flips to untied. Add tests covering INPUT-only and OUTPUT-only saves from a tied source (existing tests only cover non-tied / BOTH).

Medium

3. INPUT-only mode replays every decoder layer for nothing. With embed_quant_mode=INPUT the progress bar is layer_count + 1 and each layer hits replay_skipped_layer (full forward over all calibration data) although no downstream endpoint consumes the activations. Break out of the layer loop after the input endpoint when quant_output_embeddings is false (and lm_head is off).

4. Requantizing an already-quantized endpoint is not supported, only failing late. If the source checkpoint already has lm_head (or embed_tokens) as a quant module, the type check in module_looper.py raises NotImplementedError after model load + calibration loading. Since the script is named requant_*, either support it (dequantize → requant) or document/validate up front in optimize/requant_embed_lm_head.py.

5. Hard-coded embedding default {"bits": 8, "group_size": 32, "sym": True, "desc_act": False, "mse": 2.4} in module_looper.py duplicates (and diverges from) WeightOnlyLooper._configure_embedding_dynamic_defaults ({"bits": 8, "group_size": 32}). Share one helper so RTN and GPTQ paths get the same defaults.

6. ModelWriter now unconditionally writes model_config["tie_word_embeddings"] = runtime_config.tie_word_embeddings for all saves, not just embedding requant. This is probably fine (normal saves keep the source value), but it is a global behavior change that deserves a note in the PR description.

Low / nits

• requantize() derives quantize_config.device = DEVICE(endpoint_device.type) — drops the device index (cuda:1 → cuda).
• GPTQProcessor.pre_process_fwd_hook embedding branch: selected_out.data fails if out is not a tensor (selected_out = out fallthrough). Guard or drop the .data.
• gptq.py::_quantize_embedding: with a diagonal (one-hot) Hessian the OBQ update has no cross-column term, so this is exactly frequency-weighted RTN with frequency-driven desc_act/GAR reordering. Worth a one-line docstring so nobody expects error compensation here.
• The two github-code-quality "empty except" comments (optimize/requant_embed_lm_head.py, writer.py) are trivial to address with a comment/log.debug.
• Regarding finding 1 of the earlier automated review (tied models resolving output to the input path): for HF models get_output_embeddings() returns the distinct lm_head module (only the weight is shared), so get_module_name yields lm_head and the CLI override does apply to the output path. I don't think that finding holds as written; the real tied-model problem is item 2 above.
• New files optimize/_common.py, optimize/requant_embed_lm_head.py, and the two new test files have ruff I001 (import sort) / X | None style findings; the rest of the ruff noise on touched files is pre-existing on main.

[ZX-ModelCloud]

Follow-up remediation — all three findings from the validation review are addressed in 022994c3.

Tied embedding replacement save lifecycle. Fixed single-endpoint requantization for tied source checkpoints. Input-only saves now serialize the cloned untied output head and preserve the untied config; output-only saves locate and rewrite the tied input shard correctly. Added regression coverage for both, input-only, and output-only endpoint selection.

Unaligned embedding vocabulary packing. Fixed 2/4/8-bit embedding packing when the vocabulary size is not divisible by the packing width (for example, vocab size 17 or GPT-2 vocab size 50257). Packing now pads the tail safely, while dequantization trims the packed tail back to the original vocabulary size. Added forward and safetensors save/reload coverage.

Embedding quantization stage selection. Preserved decoder subset plans when embed_only=False, so GPTQ/calibrated embedding-plus-decoder quantization still runs the decoder quantization stages. Embedding-only mode continues to use decoder replay only. Added regression coverage for the setting propagation and subset-stage execution.

Regression suite: 100 passed, 13 skipped (CUDA-only skips). Targeted Ruff checks and git diff --check also pass.

[Qubitium]

Merge conflict

[ZX-ModelCloud]

Merge conflict

Conflict resolved.