Fix silent quantized-buffer corruption in load_quantized_model() for Qwen3.6

[dogwood-flo]

Summary

• Fixes a bug where quantized models (Qwen3.6-27B in particular) produced garbage output after a save → load round-trip.
• Added post-load validation (_check_load_state_dict_result() / _assert_quantized_modules_loaded()) so this class of bug fails fast instead of silently producing a broken model.

This pull request should be reviewed after merging (#34).

Background

The root cause is that transformers unconditionally rewrites saved checkpoint key prefixes for certain model types (e.g. Qwen3.6's qwen3_5_text: model.layers.* becomes model.language_model.layers.* on disk), while quantization_config's recorded module names stay as the original in-memory names. _replace_quantized_layers() now actually moves matched tensors to the correct key instead of only using them to infer buffer shapes, so quantized weights are no longer silently loaded as all-zero.

Details

Bug fix

• Fixed load_quantized_model() silently loading quantized layers (GPTQ/DBF) with all-zero buffers when the checkpoint's on-disk key prefixes don't match quantization_config's recorded module names (e.g. Qwen3.6: transformers unconditionally rewrites model.layers.* to model.language_model.layers.* on save, regardless of any OneComp save option). _replace_quantized_layers() now moves the matched tensors to the correct key before load_state_dict() instead of only using them to infer buffer shapes (onecomp/quantized_model_loader.py)
  • Added _check_load_state_dict_result() / _assert_quantized_modules_loaded() post-load checks that fail fast instead of silently producing a model that generates garbage
  • Generalized the generic suffix-matching fallback in _resolve_state_dict_key() to consider the full key depth instead of a hardcoded 8-component limit
• Fixed _assert_quantized_modules_loaded() raising for every DBF-quantized model: it checked for a non-existent bp attribute on DoubleBinaryLinear instead of the real scaling0/scaling2/scaling4/bp1/bp3 (onecomp/quantized_model_loader.py)
• Fixed a false-positive load failure for tied-embedding models: lm_head.weight is legitimately absent from the checkpoint (HF's save_pretrained does not duplicate a tensor sharing storage with embed_tokens.weight) and is no longer flagged as a critical missing key (onecomp/quantized_model_loader.py)

Refactor

• Extracted the tied-embeddings re-tie decision and tie_weights() call in load_quantized_model() into _retie_lm_head_if_needed() (onecomp/quantized_model_loader.py)

Test

• Added unit tests covering the above fixes (tests/onecomp/runner/test_remap_state_dict_keys.py, tests/onecomp/runner/test_load_tied_embeddings.py)

Verified

Verified the implementation with the following script:

import torch

from onecomp import (
    GPTQ,
    CalibrationConfig,
    ModelConfig,
    Runner,
    load_quantized_model,
    setup_logger,
)

setup_logger()

MODEL_ID = "Qwen/Qwen3.6-27B"
SAVE_DIR = "./Qwen3.6-27B_gptq4"

# ── 1. Quantize with GPTQ ────────────────────────────────────
model_config = ModelConfig(model_id=MODEL_ID, device="cuda:0", dtype="bfloat16")

gptq = GPTQ(
    wbits=4,
    groupsize=128,
)

runner = Runner(
    model_config=model_config,
    quantizer=gptq,
    qep=True,
    calibration_config=CalibrationConfig(max_length=512, num_calibration_samples=128),
)
# NOTE: The calibration settings above are kept compact so the demo runs
# fast and may be insufficient for real quantisation.  For higher quality,
# prefer the CalibrationConfig() defaults
# (max_length=2048, num_calibration_samples=512).
runner.run()

# ── 2. Save ───────────────────────────────────────────────────
runner.save_quantized_model(SAVE_DIR)
print(f"\nQuantized model saved to: {SAVE_DIR}")

# ── 3. Load ───────────────────────────────────────────────────
model, tokenizer = load_quantized_model(SAVE_DIR, torch_dtype=torch.bfloat16)
print(f"Loaded model type : {type(model).__name__}")
print(f"Loaded model device: {next(model.parameters()).device}")

# ── 4. Verify: simple text generation ─────────────────────────
prompt = "Fujitsu is"
inputs = tokenizer(prompt, return_tensors="pt").to(next(model.parameters()).device)

with torch.no_grad():
    output_ids = model.generate(**inputs, max_new_tokens=32, do_sample=False)

generated = tokenizer.decode(output_ids[0], skip_special_tokens=True)
print(f"\nPrompt   : {prompt}")
print(f"Generated: {generated}")