[xin3he]

… update tests

Description

some weights should kept as FP32 in tuning phase

Type of Change

New feature

Related Issues

Fixes or relates to #2332

Checklist Before Submitting

• My code has been tested locally.
• Documentation has been updated as needed.
• New or updated tests are included where applicable.
• The CUDA CI has passed. You can trigger it by commenting /azp run Unit-Test-CUDA-AutoRound.

[Copilot]

🟡 Changes recommended

Three moderate issues remain unresolved in tensor preservation, revision-aware checkpoint resolution, and model-free routing.

2 open issues

• Preserve mixed-checkpoint FP32 parameters already loaded as FP32 · New
• Preserve model-free fallback instead of raising an exception · New

What changed in this PR

Enhances mixed-precision checkpoint loading by preserving selected FP32 tensors during loading, casting, and materialization.

Changes:

• Adds FP32 tensor detection and restoration.
• Integrates preservation into regular and meta-model loading.
• Expands regression tests and updates model-free routing.

Open findings remain around FP32 detection, checkpoint revision handling, and model-free routing behavior.

File	Reviewed change
test/​unit/​common/​utils/​test_missing_tensors.py	Adds mixed-precision restoration tests.
auto_round/​utils/​model.py	Integrates FP32 retention and updates routing behavior.
auto_round/​utils/​missing_tensors.py	Detects and preserves mixed-checkpoint FP32 tensors.
auto_round/​context/​model.py	Applies FP32 retention to meta-loaded models.

_{Review effort: Lite}

[AutoRoundBot]

CI Failure Analysis Report (Unit-Test-AutoRound)

• Commit: be2e595
• Run: https://dev.azure.com/lpot-inc/neural-compressor/_build/results?buildId=78185
• Totals: 3 failed test(s), 2 cluster(s) (0 known / 2 new)

New Issues (2)

No.1 — 2 occurrence(s)

🔍 ValueError: model_free=True does not support static KV or attention quantization.

📝 Basic info

• Affected tests (2): test_model_free_route_rejects_static_attention_quantization[True-static_kv_dtype], test_model_free_route_rejects_static_attention_quantization[True-static_attention_dtype]
• Logs (1): unittest_test_model_free_utils.log

🖥️ Log excerpt

unit/test_cpu/utils/test_model_free_utils.py:1284: in test_model_free_route_rejects_static_attention_quantization
    assert not is_model_free_route("model", "W4A16", 0, True, kwargs)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
../auto_round/utils/model.py:3216: in is_model_free_route
    raise ValueError("model_free=True does not support static KV or attention quantization.")
E   ValueError: model_free=True does not support static KV or attention quantization.

✨ AI analysis

• Category: Code Regression

• Confidence: high

• Root cause: The PR changed is_model_free_route so explicit model_free=True raises when static_kv_dtype or static_attention_dtype is set, whereas the existing contract and test expect the function to return False and use the regular quantization route.

• Suggested fix: Evaluate the static KV/attention quantization guard before the explicit model_free branch, returning False instead of raising.

• Patch:

  diff --git a/auto_round/utils/model.py b/auto_round/utils/model.py
  --- a/auto_round/utils/model.py
  +++ b/auto_round/utils/model.py
  @@ -3211,12 +3211,10 @@ def is_model_free_route(
       has_static_attention_quantization = any(
           kwargs.get(key) is not None for key in ("static_kv_dtype", "static_attention_dtype")
       )
  -    if explicit:
  -        if has_static_attention_quantization:
  -            raise ValueError("model_free=True does not support static KV or attention quantization.")
  -        return True
       if has_static_attention_quantization:
           return False
  +    if explicit:
  +        return True
       # Only auto-route when format is auto_round (or not specified).

No.2 — 1 occurrence(s)

🔍 assert torch.bfloat16 == torch.float32

📝 Basic info

• Affected tests (1): test_fake_format_restores_fp32_and_copies_mtp_tensors
• Logs (1): unittest_test_fake_format.log

🖥️ Log excerpt

unit/test_cpu/export/test_fake_format.py:448: in test_fake_format_restores_fp32_and_copies_mtp_tensors
    assert restored.dtype == torch.float32
E   assert torch.bfloat16 == torch.float32
E    +  where torch.bfloat16 = tensor([ 1., -2.,  3.,  4.], dtype=torch.bfloat16).dtype
E    +  and   torch.float32 = torch.float32

✨ AI analysis

• Category: Test Case Issue

• Confidence: high

• Root cause: The FP32 restoration logic intentionally skips checkpoints whose source tensors are all FP32. This fixture creates both the special norm tensor and the MTP tensor as FP32, so the mixed-precision detection is false and the norm remains BF16; the MTP copy itself succeeds.

• Suggested fix: Make the synthetic MTP tensor BF16 so the fixture represents the intended mixed-precision checkpoint and allows the FP32 norm tensor to be restored.

• Patch:

  --- a/test/unit/test_cpu/export/test_fake_format.py
  +++ b/test/unit/test_cpu/export/test_fake_format.py
  @@ -412,7 +412,7 @@ def test_fake_format_restores_fp32_and_copies_mtp_tensors(tmp_path, monkeypatch):
       source_dir = str(tmp_path / "source")
       os.makedirs(source_dir)
       norm_weight = torch.tensor([1.0001, -2.0002, 3.0003, 4.0004], dtype=torch.float32)
  -    mtp_weight = torch.randn(8, 4)
  +    mtp_weight = torch.randn(8, 4, dtype=torch.bfloat16)
       save_file(
           {
               "model.language_model.layers.0.linear_attn.norm.weight": norm_weight,

Notes

• Only top-3 issues receive AI analysis.
• To request an additional fix from Copilot, use "Quote reply" on the PR comment and @mention Copilot.