mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-11-10 09:34:01 +08:00
realAsma
c4b662fbc8
## What does this PR do? **Type of change:** Bug fix **Overview:** Fix `AttributeError: Can't get local object 'add_hook_to_module.<locals>.new_forward'` when saving a quantized model a second time after restoring it with `device_map="auto"`. When a model is loaded with `device_map="auto"`, accelerate's `add_hook_to_module` patches every submodule (including `TensorQuantizer` instances) and injects three instance attributes: `_hf_hook`, `_old_forward`, and `forward` (a `functools.partial` wrapping a local function). These are not picklable and were leaking into the modelopt state dict collected by `get_modelopt_state()`, causing `torch.save` to fail. This PR adds the three accelerate-injected attributes to `TensorQuantizer._skip_properties_for_save_restore` so they are excluded from the serialized state, matching the existing pattern used for HuggingFace and DeepSpeed attributes. ## Usage ```python mto.enable_huggingface_checkpointing() # Quantize and save model = AutoModelForCausalLM.from_pretrained(name, device_map="auto") model = mtq.quantize(model, mtq.FP8_DEFAULT_CFG, forward_loop=forward_loop) model.save_pretrained(save_dir) # Restore and save again (this previously failed) model2 = AutoModelForCausalLM.from_pretrained(save_dir, device_map="auto") model2.save_pretrained(save_dir_round2) # now works ``` ## Testing - Added unit test `test_tensor_quantizer_modelopt_state_with_accelerate_hook` in `tests/unit/torch/quantization/plugins/test_accelerate.py` that verifies accelerate hook attributes are excluded from modelopt state and the state dict remains picklable. ## Before your PR is "*Ready for review*" - **Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)** and your commits are signed. - **Is this change backward compatible?**: Yes — only adds entries to a skip set; existing saved checkpoints are unaffected. - **Did you write any new necessary tests?**: Yes - **Did you add or update any necessary documentation?**: No (internal fix, no API change) - **Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**: No ## Additional Information The root cause is in accelerate's `add_hook_to_module`, which defines `new_forward` as a local function and binds it via `functools.partial` onto `module.forward`. Since local functions cannot be pickled, any `TensorQuantizer` that has been hooked by accelerate becomes unserializable unless these attributes are excluded. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Enhanced compatibility with accelerate library by excluding framework-specific hooks and attributes from model state serialization, preventing issues during save/restore operations. * **Tests** * Added test to validate that accelerate-related attributes are properly excluded from model state and that the state remains picklable. * **Public API** * TensorQuantizer is now publicly exported. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: realAsma <akuriparambi@nvidia.com>
83 lines
3.0 KiB
Python
83 lines
3.0 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import pickle
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
import modelopt.torch.quantization as mtq
|
|
from modelopt.torch.quantization.nn import QuantLinearConvBase, TensorQuantizer
|
|
|
|
try:
|
|
from accelerate.hooks import ModelHook, add_hook_to_module
|
|
except ImportError:
|
|
pytest.skip("accelerate not available", allow_module_level=True)
|
|
|
|
|
|
def test_linear_with_accelerate_monkey_patched_forward():
|
|
module_test = nn.Linear(16, 16)
|
|
add_hook_to_module(module_test, ModelHook())
|
|
|
|
mtq.replace_quant_module(module_test)
|
|
assert module_test._old_forward.__func__ == QuantLinearConvBase.forward
|
|
|
|
module_test.input_quantizer.enable_calib()
|
|
module_test.weight_quantizer.enable_calib()
|
|
|
|
module_ref = nn.Linear(16, 16)
|
|
mtq.replace_quant_module(module_ref)
|
|
|
|
module_ref.load_state_dict(module_test.state_dict())
|
|
|
|
x = torch.randn(1, 16)
|
|
out1 = module_test(x)
|
|
out2 = module_ref(x)
|
|
assert torch.allclose(out1, out2)
|
|
|
|
module_test.input_quantizer.load_calib_amax()
|
|
module_test.weight_quantizer.load_calib_amax()
|
|
|
|
assert module_test.input_quantizer.amax is not None
|
|
assert module_test.weight_quantizer.amax is not None
|
|
|
|
|
|
def test_tensor_quantizer_modelopt_state_with_accelerate_hook():
|
|
"""Verify accelerate hook attributes are excluded from modelopt state.
|
|
|
|
When accelerate's add_hook_to_module patches a TensorQuantizer, it adds
|
|
_hf_hook, _old_forward, and an instance-level forward (a functools.partial
|
|
wrapping a local function). These must be excluded from the modelopt state
|
|
dict, otherwise torch.save / pickle will fail with:
|
|
AttributeError: Can't get local object 'add_hook_to_module.<locals>.new_forward'
|
|
"""
|
|
tq = TensorQuantizer()
|
|
add_hook_to_module(tq, ModelHook())
|
|
|
|
# The hook should have injected these instance attributes
|
|
assert hasattr(tq, "_hf_hook")
|
|
assert hasattr(tq, "_old_forward")
|
|
assert "forward" in tq.__dict__
|
|
|
|
# None of the accelerate attributes should appear in the modelopt state
|
|
state = tq.get_modelopt_state()
|
|
accelerate_attrs = {"_hf_hook", "_old_forward", "forward"}
|
|
leaked = accelerate_attrs & state.keys()
|
|
assert not leaked, f"Accelerate attributes leaked into modelopt state: {leaked}"
|
|
|
|
# The state dict must be picklable (torch.save uses pickle internally)
|
|
pickle.dumps(state)
|