Make FqnToConfig handle module swap configs #3492
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3492
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit 264b3c2 with merge base f3342a0 ( BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
bot
added
the
CLA Signed
andrewor14
force-pushed
the
fqn-to-config-qat
branch
from
50b55f5 to
8bfc239
Compare
vkuzo
reviewed
Dec 16, 2025
torchao/quantization/quant_api.py
Outdated
| module.to(device=device) | ||
| return | ||
| if isinstance(config, AOBaseConfig): | ||
| _fqn_to_config_replace(config, model, device=device) |
Contributor
There was a problem hiding this comment.
this is hard to follow, would it be cleaner to use iteration instead of recursion here
jcaip
reviewed
Dec 16, 2025
| or ("_default" in config.fqn_to_config and _is_linear(module)) | ||
| ): | ||
| # this replaces inplace, so no need to reassign | ||
| _fqn_to_config_handler(module, module_fqn, config) |
Contributor
There was a problem hiding this comment.
Maybe we can implement something like
def set_module_by_fqn(model, fqn, new_module):
# set model[fqn] to new_moduleso then we can keep it iterative? Only thing is that we have to get the modules in top-down order (this does appear to be how named_modules is implemented)
for module_fqn, module in model.named_modules():
if (
fqn_matches_fqn_config(module_fqn, config)
or _module_param_matches_fqn_config(module, module_fqn, config)
or ("_default" in config.fqn_to_config and _is_linear(module))
):
new_module = _fqn_to_config_handler(module, module_fqn, config)
# reassign module
set_module_by_fqn(model, fqn, new_module)
andrewor14
force-pushed
the
fqn-to-config-qat
branch
from
8bfc239 to
f8af9cd
Compare
**Summary:** `FqnToConfig` does not handle module swap configs like `QATConfig` today because it does not reassign the replaced modules, but instead assumes there is no need to do so because the underlying tensors are swapped. This breaks when the user tries to use `FqnToConfig` with `QATConfig`, which does not rely on tensor subclasses. This commit fixes this by making the replacement logic in `FqnToConfig` more general. Fixes #3490. **Test Plan:** ``` python test/quantization/test_quant_api.py -k test_fqn_config_quantized_nested_module_module_swap ```
andrewor14
force-pushed
the
fqn-to-config-qat
branch
from
f8af9cd to
264b3c2
Compare
Contributor
Author
vkuzo
approved these changes
Dec 17, 2025
jcaip
reviewed
Dec 17, 2025
| child_name = module_fqn.split(".")[-1] | ||
| parent_fqn = module_fqn.removesuffix(child_name).removesuffix(".") | ||
| parent_module = named_modules[parent_fqn] | ||
| setattr(parent_module, child_name, replacement) |
Contributor
There was a problem hiding this comment.
I think we're missing a return here at the end
Contributor
Author
There was a problem hiding this comment.
I changed the next case to elif, is that OK?
jcaip
approved these changes
Dec 17, 2025
Contributor
Author
|
Test failure doesn't seem related, merging this. Thanks |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary:
FqnToConfigdoes not handle module swap configs likeQATConfigtoday because it does not reassign the replaced modules, but instead assumes there is no need to do so because the underlying tensors are swapped. This breaks when the user tries to useFqnToConfigwithQATConfig, which does not rely on tensor subclasses. This commit fixes this by making the replacement logic inFqnToConfigmore general.Fixes #3490.
Test Plan: