Autoquant v2 initial version #1240
Conversation
Summary: We refactored the v1 to do benchmark for subgraphs of (prev_op -> linear -> post_op) in order to get more accurate estimation of timing. One issue here is now we need to care about batch size of the subgraph, so we'd need the batch size dimension to use symbolic shape, seems that it does not have good support on torch.compile right now More improvements: * current batch size adjustment code is hardcoded to work for llama model, need to think of a way to generalize it * using canonicalized subgraph as key for the cache to reduce the number of times we need to do benchmarking * add accuracy sanity checks Test Plan: Testing with torchao/_models/llama/generate.py ``` python generate.py --checkpoint_path $CHECKPOINT_PATH/$MODEL_REPO/model.pth --compile --compile_prefill --quantization autoquant_v2-int4 ``` Reviewers: Subscribers: Tasks: Tags:
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1240
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
| torch.nn.Linear(*new_shape, dtype=weight_val.dtype), | ||
| ).cuda() | ||
|
|
||
| else: |
There was a problem hiding this comment.
this file has some complexity for extracting (prev_op -> linear1 -> maybe_linear_2 -> next_ops) because the models we originally studied had back to back linears. If you only care about transformer models, you can simplify this code quite a bit by removing the special logic for extraction of the second linear. Happy to point to the right places in the code if needed.
| return True | ||
| return False | ||
|
|
||
| def debug_single_linear( |
There was a problem hiding this comment.
depending on what you're using this file for, this function also might be deleteable
There was a problem hiding this comment.
yeah will refine more when it's closer to land, right now just experimenting to see if this approach helps improve things over the original approach on models we care about
Summary:
We refactored the v1 to do benchmark for subgraphs of (prev_op -> linear -> post_op) in order to get more accurate estimation of timing. One issue here is now we need to care about batch size of the subgraph, so we'd need the batch size dimension to use symbolic shape, seems that it does not have good support on torch.compile right now
More improvements:
Issues
Test Plan:
Testing with torchao/_models/llama/generate.py
Reviewers:
Subscribers:
Tasks:
Tags: