cleaned up and tested tp support #976
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/976
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 273fd8c with merge base c0a81f9 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
torchao/_models/llama/generate.py
Outdated
| @@ -30,7 +31,7 @@ def device_sync(device): | |||
| wd = Path(__file__).parent.parent.resolve() | |||
| sys.path.append(str(wd)) | |||
|
|
|||
| from torchao._models.llama.model import Transformer, prepare_inputs_for_model | |||
| from model import Transformer, prepare_inputs_for_model | |||
There was a problem hiding this comment.
i think we want this to be the same as before, otherwise it can accidentally pull in gpt-fast model and it becomes annoying to test/debug
| funcol = None | ||
|
|
||
| from torchao._models.llama.model import Attention, FeedForward, Transformer | ||
| from torchao.quantization.GPTQ import WeightOnlyInt4Linear |
There was a problem hiding this comment.
we are not usingWeightOnlyInt4Linear any more in torchao I think, is this just for GPTQ?
|
@debajyotidatta the way we are supporting tp by composing quantized tensor subclass and DTensor: #939, I believe the way its done in gpt-fast is outdated, would you be interested in integrating the new tp flow in generate.py instead? |
Sounds good! I will do that. |
No description provided.