-
Notifications
You must be signed in to change notification settings - Fork 14
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
Hierarchical Test template for support Grok/GPT via PAXML #141
Conversation
Approved. Let's see if @amaslenn has any comments and wait for @srinivas212 to merge this. @amaslenn, this PR supports two different subtests in the JAX Toolbox test templates. This allows us to support Grok and GPT simultaneously in a hierarchical manner. The related PR is https://github.com/Mellanox/cloudaix/pull/2. |
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.
This PR comes with zero tests, yet changes a lot. Please add unit tests.
Summary
Modification to Jaxtoolbox to support hierarchical test template to support both GPT and Grok models via PAXML.
Test Plan
Internal testing and pytests
Additional Notes
Needs additional testing on same container. XLA Flags seems to have changed. It requires debugging to make Grok working on the same container.