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.
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Sweep across KV cache layouts #662
base: main
Are you sure you want to change the base?
Sweep across KV cache layouts #662
Changes from 65 commits
91f4398
18ba546
e78faee
6c04467
b839d63
baaeda7
f0588ee
d85990b
cbb90bf
929ad8d
f3cbd3e
0a6de55
f907515
8333659
8c0707d
674ba29
95f846f
858f4e3
1c20d61
a9b381a
ec1d76e
63edcc2
268ee64
c1490f6
0fe1d51
554d014
6458a89
76bda5d
42bd412
df5c66e
c44da0c
d84f4c9
185563d
4291c4c
c638c63
128a691
8387cd7
b422007
f69f405
c37ad18
4b3be09
65c2289
d5637e4
9c0d6a5
58405b1
187cb3d
958938c
e37e495
a7fb24b
cb2db18
7311225
000e935
ceda588
69bfc92
ca107c5
1d417e4
0c9a6f3
296119b
58fec1e
74a1e70
70dd34c
bd29b6e
65e12a4
746bbda
3d0214e
0151bb0
951200e
22cae3b
081891f
0be7296
b6fd256
220cfd2
4b9c8a3
7fd12c6
7c9ccae
5d798a6
4f0a12c
File filter
Filter by extension
Conversations
Jump to
There are no files selected for viewing
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You can skip this since there is already
in the https://github.com/GoogleCloudPlatform/ml-auto-solutions/pull/288/files#diff-0ca160e04db7871a01e5f1aafc9705f3f52efe9b1321b2901289554e806f55d9R113
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
My understanding is that
write_config_raw_keys_for_gcs
is called once duringinitialize
(https://github.com/google/maxtext/blob/main/MaxText/pyconfig.py#L225), and writes the default values of prefill and ar axis orders to GCS. And so each time we loop for a different prefill/ar axis, we need to explicitly callmax_utils.write_config_raw_keys_for_gcs
again to make sure we write the updated values to GCS.Did you take this out in your code, and it worked for you?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@morgandu should I do a test without
max_utils.write_config_raw_keys_for_gcs
?