Fix init_quantized_weights: model kwargs leak into dispatch, missing tie_weights() - #2161
Conversation
…tie_weights Two defects on the init_quantized_weights path (also reached via hf_ptq.py --low_memory_mode): 1. patched_from_pretrained forwarded **kwargs verbatim into load_checkpoint_and_dispatch(), so any model-construction kwarg raised TypeError. attn_implementation is the common case: 'load_checkpoint_and_dispatch() got an unexpected keyword argument attn_implementation'. It now goes to cls.from_config(), where it belongs. 2. tie_weights() was never called before quantization. Tied parameters such as lm_head.weight (tie_word_embeddings=true, e.g. Qwen2.5-0.5B) are not present in the checkpoint, so they stayed on meta and dispatch_model() raised 'Cannot copy out of meta tensor; no data!'. accelerate documents tie_weights() as a prerequisite of load_checkpoint_and_dispatch(). Reproduced with Qwen/Qwen2.5-0.5B-Instruct, nvfp4, --low_memory_mode on modelopt 0.43.0; both failures occur before any weight is written. Note: these two fixes let the path run to completion but do not make its output correct - see NVIDIA#2160 for the separate root-cause defect (quantization runs on meta tensors before real weights are loaded). Signed-off-by: spped2000 <spped2000@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: spped2000 <spped2000@gmail.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthrough
ChangesAccelerate loading
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelopt/torch/quantization/plugins/accelerate.py`:
- Around line 225-237: Update the model-construction kwargs handling around
cls.from_config to pop both dtype and torch_dtype from kwargs before
load_checkpoint_and_dispatch receives them, while preserving the existing
precedence and fallback selection. Pass the normalized selected value explicitly
as dtype= when constructing the model, and ensure neither alias remains in
kwargs.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 52088e16-a681-425e-86ca-e488c54f6b0d
📒 Files selected for processing (1)
modelopt/torch/quantization/plugins/accelerate.py
…citly CodeRabbit correctly noted that dtype/torch_dtype were read with kwargs.get() and therefore stayed in kwargs. load_checkpoint_and_dispatch() accepts dtype but NOT torch_dtype, so a caller using the legacy alias hit the same TypeError class this PR fixes for attn_implementation. Both aliases are now popped, precedence (dtype > torch_dtype > config) is unchanged, and the resolved value is passed explicitly to load_checkpoint_and_dispatch() so weight casting behaves exactly as before for callers who passed dtype. Signed-off-by: spped2000 <spped2000@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: spped2000 <spped2000@gmail.com>
|
Good catch, fixed and pushed. Confirmed the concern is real: Both aliases are now popped, precedence ( |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
modelopt/torch/quantization/plugins/accelerate.py (1)
236-266: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winUse
config.dtypefor the fallback.When
config.dtypeisNone,config.torch_dtypealso returnsNonebecause it is a deprecated alias. The model is then constructed with the defaultfloat32, andload_checkpoint_and_dispatch()receivesdtype=None. This bypasses the intendedfloat16fallback and can cause low-memory loading failures.Use
getattr(config, "dtype", None) or torch.float16and add a regression test for a config without a dtype.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/quantization/plugins/accelerate.py` around lines 236 - 266, Update the torch_dtype fallback in the model construction flow to use config.dtype, defaulting to torch.float16 when it is missing or None; do not use the deprecated config.torch_dtype alias. Ensure the resolved dtype is passed consistently to both cls.from_config and load_checkpoint_and_dispatch, and add a regression test covering a configuration without a dtype.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@modelopt/torch/quantization/plugins/accelerate.py`:
- Around line 236-266: Update the torch_dtype fallback in the model construction
flow to use config.dtype, defaulting to torch.float16 when it is missing or
None; do not use the deprecated config.torch_dtype alias. Ensure the resolved
dtype is passed consistently to both cls.from_config and
load_checkpoint_and_dispatch, and add a regression test covering a configuration
without a dtype.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e4aab36c-0c2b-49d9-8a1d-2cd43f0eec59
📒 Files selected for processing (1)
modelopt/torch/quantization/plugins/accelerate.py
…regression test getattr(config, "torch_dtype", torch.float16) never reaches its default on transformers >= 5: torch_dtype is a deprecated alias of dtype that RETURNS None when unset rather than being absent, so the attribute exists and getattr hands back None. The model was then built as float32 and dtype=None reached load_checkpoint_and_dispatch, on the code path whose entire purpose is loading under tight memory. Verified on transformers 5.13.1: PretrainedConfig().torch_dtype -> None getattr(cfg, 'torch_dtype', torch.float16) -> None # not float16 Reads config.dtype first, falls back to the alias, then to float16 explicitly. Adds a unit test covering kwarg precedence, kwarg consumption, and the no-dtype-in-config fallback. Signed-off-by: spped2000 <spped2000@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: spped2000 <spped2000@gmail.com>
|
Third finding addressed (the out-of-diff one) — and it reproduces, so thank you for catching it. Verified on transformers 5.13.1: >>> PretrainedConfig().torch_dtype
None
>>> getattr(PretrainedConfig(), "torch_dtype", torch.float16)
None # not float16 — the getattr default never firesBecause Now reads Note the PR now carries three distinct fixes. Happy to split the dtype one out if you would rather review them separately. |
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unit/torch/quantization/plugins/test_accelerate.py`:
- Line 94: At module scope in the test module, guard the optional dependency
with pytest.importorskip("transformers"), then move the PretrainedConfig import
into the top-level module imports alongside the other imports. Remove any
deferred or test-local PretrainedConfig import while preserving the existing
test behavior.
- Around line 85-93: Update test_init_quantized_weights_dtype_resolution to
invoke the real patched_from_pretrained loader with checkpoint I/O stubbed,
rather than testing a local resolve() helper. Configure conflicting dtype
sources and a non-None deprecated config.torch_dtype, then assert
init_quantized_weights/from_config receives the selected dtype and
attn_implementation, while load_checkpoint_and_dispatch receives only the
supported dispatch arguments; cover the precedence dtype > torch_dtype >
config.dtype > config.torch_dtype > torch.float16.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 172afdc2-0dc2-4fc0-ac70-80e4040120ad
📒 Files selected for processing (2)
modelopt/torch/quantization/plugins/accelerate.pytests/unit/torch/quantization/plugins/test_accelerate.py
🚧 Files skipped from review as they are similar to previous changes (1)
- modelopt/torch/quantization/plugins/accelerate.py
| def test_init_quantized_weights_dtype_resolution(): | ||
| """dtype/torch_dtype must not leak into load_checkpoint_and_dispatch(). | ||
|
|
||
| Both are model-construction kwargs: `load_checkpoint_and_dispatch()` | ||
| accepts `dtype` but not `torch_dtype`, so forwarding kwargs verbatim | ||
| raised TypeError for callers using the legacy alias. The fallback also has | ||
| to survive `config.torch_dtype` being a deprecated alias that returns None | ||
| instead of being absent, which defeats a `getattr(..., default)` fallback. | ||
| """ |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline tests/unit/torch/quantization/plugins/test_accelerate.py --items all
rg -n -C 6 \
'test_init_quantized_weights_dtype_resolution|def resolve|patched_from_pretrained|init_quantized_weights|load_checkpoint_and_dispatch|attn_implementation' \
tests/unit/torch/quantization/plugins/test_accelerate.py \
modelopt/torch/quantization/plugins/accelerate.pyRepository: NVIDIA/Model-Optimizer
Length of output: 16386
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- test file ---'
sed -n '1,155p' tests/unit/torch/quantization/plugins/test_accelerate.py
printf '%s\n' '--- production loader ---'
sed -n '212,280p' modelopt/torch/quantization/plugins/accelerate.py
printf '%s\n' '--- static checks ---'
python3 - <<'PY'
from pathlib import Path
import ast
test_path = Path("tests/unit/torch/quantization/plugins/test_accelerate.py")
source = test_path.read_text()
tree = ast.parse(source)
target = next(
node for node in tree.body
if isinstance(node, ast.FunctionDef)
and node.name == "test_init_quantized_weights_dtype_resolution"
)
calls = []
assertions = []
dict_literals = []
for node in ast.walk(target):
if isinstance(node, ast.Call):
if isinstance(node.func, ast.Name):
calls.append(node.func.id)
elif isinstance(node.func, ast.Attribute):
calls.append(node.func.attr)
elif isinstance(node, ast.Assert):
assertions.append(ast.unparse(node.test))
elif isinstance(node, ast.Dict):
keys = [
ast.literal_eval(k) for k in node.keys
if isinstance(k, (ast.Constant, ast.Str))
]
if keys:
dict_literals.append(keys)
print("calls:", sorted(set(calls)))
print("assertions:")
for item in assertions:
print(" ", item)
print("dict keys:", dict_literals)
print("uses_real_loader:", any(
name in {"patched_from_pretrained", "init_quantized_weights",
"load_checkpoint_and_dispatch"}
for name in calls
))
print("asserts_attn_implementation:", any(
"attn_implementation" in item for item in assertions
))
PYRepository: NVIDIA/Model-Optimizer
Length of output: 9068
Exercise the real patched loader in this regression test.
resolve() only mirrors production. It does not call patched_from_pretrained, init_quantized_weights, or load_checkpoint_and_dispatch, and it does not assert attn_implementation. Invoke the real loader with checkpoint I/O stubbed, then assert the from_config() and dispatch arguments. Use conflicting values to test dtype > torch_dtype > config.dtype > config.torch_dtype > torch.float16, including a non-None deprecated config.torch_dtype.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unit/torch/quantization/plugins/test_accelerate.py` around lines 85 -
93, Update test_init_quantized_weights_dtype_resolution to invoke the real
patched_from_pretrained loader with checkpoint I/O stubbed, rather than testing
a local resolve() helper. Configure conflicting dtype sources and a non-None
deprecated config.torch_dtype, then assert init_quantized_weights/from_config
receives the selected dtype and attn_implementation, while
load_checkpoint_and_dispatch receives only the supported dispatch arguments;
cover the precedence dtype > torch_dtype > config.dtype > config.torch_dtype >
torch.float16.
Source: Path instructions
The previous test asserted against a local mirror of the resolution logic, so it would have passed even if patched_from_pretrained were broken. It now drives the real init_quantized_weights context manager with checkpoint I/O stubbed (mtq.quantize/compress, the accelerate device-map helpers and load_checkpoint_and_dispatch are patched) and asserts what each callee actually receives: - from_config gets the resolved dtype and attn_implementation; - dispatch receives neither torch_dtype nor attn_implementation, which it cannot accept, but does receive the resolved dtype; - tie_weights() is called before quantization; - with no dtype anywhere the float16 fallback fires, including the case where the deprecated config.torch_dtype alias returns None. Verified it fails against the pre-PR implementation and passes with the fix. Also moved the transformers import to module scope behind pytest.importorskip, per review. Signed-off-by: spped2000 <spped2000@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: spped2000 <spped2000@gmail.com>
|
Both review points addressed — the first one was a fair hit and I have rewritten the test. The test was testing a copy of the logic, not the code. As written it mirrored the resolution order in a local
I checked it actually protects the code: reverting
(One note on the suggested precedence |
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unit/torch/quantization/plugins/test_accelerate.py`:
- Line 98: Move the accel_plugin import from the test method to module scope,
placing it after the module-level pytest.importorskip("transformers") guard.
Keep the existing test behavior unchanged; only retain a deferred import if an
optional-dependency constraint requires it, and document that reason briefly.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 4c7fe034-d661-4717-b8eb-1ba87daae621
📒 Files selected for processing (1)
tests/unit/torch/quantization/plugins/test_accelerate.py
Per the repo's coding guidelines, imports belong at the top of test files. The plugin import now sits after the module-level accelerate/transformers guards. Signed-off-by: spped2000 <spped2000@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: spped2000 <spped2000@gmail.com>
|
Done — That is all five review points from this round addressed. Summary of where the PR stands: three independent fixes on the |
1 participant
Two small defects on the
init_quantized_weightspath (public API; also whathf_ptq.py --low_memory_modeuses). Both fire before any weight is written, so the path cannot complete at all on common models.1. Model-construction kwargs reach
load_checkpoint_and_dispatch()patched_from_pretrainedforwards**kwargsverbatim, soattn_implementation(documented inhf_ptq.py's own CLI) raises:It now goes to
cls.from_config(), where a construction kwarg belongs.2.
tie_weights()is never called before quantizationTied parameters such as
lm_head.weight(tie_word_embeddings: true— e.g. Qwen2.5-0.5B) are absent from the checkpoint, so they stay on meta anddispatch_model()raises:accelerate documents
tie_weights()as a prerequisite ofload_checkpoint_and_dispatch().Reproduction
Qwen/Qwen2.5-0.5B-Instruct(local dir),--qformat nvfp4 --low_memory_mode, modelopt 0.43.0, NGCnvcr.io/nvidia/vllm:26.05.post1-py3, GB10/SM121 aarch64. Failure 1 fires immediately; with it patched, failure 2 fires at dispatch.Scope — please read alongside #2160
These two fixes let the path run to completion, but the resulting checkpoint is still numerically wrong: quantization and compression execute on
init_empty_weights()meta tensors before real weights load, giving half-sizedweight_scaleand dequant cosine 0.756 vs the BF16 source. That root cause is filed separately as #2160 and is not addressed here — I kept this PR to the two mechanical bugs so it can be reviewed independently.Disclosure: prepared with assistance from Claude (Anthropic); all failures above were reproduced on real hardware before writing the patch.
Summary by CodeRabbit