Varying PartitionSpec of FlaxModel doesn't change IPS #1802
Unanswered
mattiasmar
asked this question in
Q&A
Replies: 0 comments
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
-
Reposting this question question here.
This Collab attempts to parallelize a very simple flax MLP model using pjit.
The problem/question is why varying the PartitionSpec in this example hardly has any effect on the time performance.
Any ideas/insights would be very welcomed!
Beta Was this translation helpful? Give feedback.
All reactions