[WIP] Make AWQ more general #2400
Conversation
Summary: * Added AWQConfig that takes a base config and made corresponding changes in other parts of the flow Test Plan: TODO Reviewers: Subscribers: Tasks: Tags:
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2400
Note: Links to docs will display an error until the docs builds have been completed. ❌ 2 New Failures, 1 Unrelated FailureAs of commit d682cb5 with merge base e4f2715 ( NEW FAILURES - The following jobs have failed:
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. |
facebook-github-bot
added
the
CLA Signed
| eps=eps, | ||
| ) | ||
| else: | ||
| observer = AWQObserver2( |
There was a problem hiding this comment.
can you not add kwargs to the AWQObserver and just check 'base_config' in kwargs?
There was a problem hiding this comment.
yes, this is temporary, I think we can deprecate the old one in the end
|
|
||
|
|
||
| @dataclass | ||
| class AWQConfig(AOBaseConfig): |
There was a problem hiding this comment.
Ok this is consolidating with quantize_ api's config based design?
| dummy_mod = DummyModule(observed_linear.weight * equalization_scale) | ||
| quant_mod = base_config_handler(dummy_mod, config.base_config) |
There was a problem hiding this comment.
I am not sure whats happening here?. Isnt module already nn.Module?
There was a problem hiding this comment.
this is just trying to quantize the weight with the quantization type specified by config.base_config
| if config.set_inductor_config: | ||
| torchao.quantization.utils.recommended_inductor_config_setter() | ||
|
|
||
| observed_linear = module |
There was a problem hiding this comment.
If this is for linear only should you not assert that this nn.Linear? Plus how to you make sure this function is called only on nn.Linear?
There was a problem hiding this comment.
yeah that's true, will add an assert, we rely on user to use quantize_ correctly (it's through specifying the filter_fn arg in quantize_ API)
ao/torchao/quantization/quant_api.py
Line 578 in 4e3d019
Summary:
Test Plan:
TODO
Reviewers:
Subscribers:
Tasks:
Tags: