Add a CPU check and CI (#15) #25
Draft
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
tengyifei
reviewed
Jan 10, 2025
.github/workflows/cpu_test.yml
Outdated
| @@ -0,0 +1,29 @@ | |||
| name: CPU tests | |||
There was a problem hiding this comment.
looks alright but need rebase
tengyifei
reviewed
Jan 10, 2025
|
|
||
| @functools.partial( | ||
| jax.jit, | ||
| out_shardings=name_to_sharding, |
There was a problem hiding this comment.
Did you test this on 405B? When I did something similar locally, the jit worked for 8B but the graph used >1TiB of HBM on 405B on two slices. The problem I encountered was that each device was still trying to generate the entire weight.
Similar to #13, we also add a CPU github action. This action will run `pytest` on the repo. Currently there is only one test, which is the Llama test in torch_xla_models. In order to run the test today, we need a HF_TOKEN. I created a personal read only token and #14 tracks avoiding the need for HF_TOKEN, after which I'll need to remember to invalidate the token.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Similar to #13, we
also add a CPU github action. This action will run
pyteston the repo.Currently there is only one test, which is the Llama test in
torch_xla_models.
In order to run the test today, we need a HF_TOKEN. I created a personal
read only token and
#14 tracks avoiding
the need for HF_TOKEN, after which I'll need to remember to invalidate
the token.