Skip to content

Do not cache input args in dynamo bridge#6553

Merged
JackCaoG merged 6 commits intomasterfrom
lsiyuan/dynamo-cache-no-input
Feb 21, 2024
Merged

Do not cache input args in dynamo bridge#6553
JackCaoG merged 6 commits intomasterfrom
lsiyuan/dynamo-cache-no-input

Conversation

@lsy323
Copy link
Copy Markdown
Collaborator

@lsy323 lsy323 commented Feb 16, 2024

Do not cache FX Graph args in GraphInputMatcher. It's not necessary to cache those tensors, because user inputs will be used in each run.

Without this change, the cached input tensors will cause more device memory usage as new graph is cached.

xla_args is the user inputs passed down from Dynamo, the cached tensors are released based on the tensor_id of xla_args, when each GraphInputMatcher is constructed.

@lsy323 lsy323 requested review from JackCaoG and wonjoo-wj February 21, 2024 00:48
@lsy323 lsy323 marked this pull request as ready for review February 21, 2024 00:48
@lsy323
Copy link
Copy Markdown
Collaborator Author

lsy323 commented Feb 21, 2024

Hi @wonjoolee95 @JackCaoG, this PR is ready for review. The CI failed because of the staled GCP token. I ran the dynamo tests locally and they passed.

Comment thread test/dynamo/test_num_output.py
Comment thread test/dynamo/test_graph_input_matcher.py
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants