Add Pathways Recipe Support for Scale Testing #1220

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.

Sign up for GitHub

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

Jump to bottom

Closed

SujeethJinesh wants to merge 1 commit into main from sujinesh/add_scale_testing_scripts

Collaborator

SujeethJinesh commented Jan 31, 2025 •

edited

Loading

Description

This PR contains the following changes:

Adds the following 2 recipes:

a. Benchmarking Pathways and McJAX code
b. Long running test on Pathways.

This is done to improve code/recipe sharing between developers during testing and make benchmarking simpler by
sharing a single, well-defined config. This will be used for testing and verification on v6e capacity.
By default allows all configs to run pathways (enable_single_controller, disable zarr3, enable_pathways_goodput etc.), with additional configs that can be added by the user with pathways_tuning_params.
Allow removing problematic XLA flags for pathways with pathways_xla_flag_options or subsequently adding xla flags to different pods.
Cleans up a lot of old pathways specific configs and adds some new ones as well.

Tests

Run multiple rounds on multiple v6e clusters

Checklist

Before submitting this PR, please make sure (put X in square brackets):

I have performed a self-review of my code.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run end-to-end tests tests and provided workload links above if applicable.
I have made or will make corresponding changes to the doc if needed.

SujeethJinesh changed the title ~~Add Pathways Benchmarking Recipes for Scale Testing~~ Add Pathways Benchmarking Recipes for Scale Testing (and fix regression from #1094)

SujeethJinesh force-pushed the sujinesh/add_scale_testing_scripts branch from 38aed8e to 4619286 Compare

January 31, 2025 07:32

SujeethJinesh changed the title ~~Add Pathways Benchmarking Recipes for Scale Testing (and fix regression from #1094)~~ Add Pathways Benchmarking Recipes for Scale Testing (and fix bugs)

SujeethJinesh force-pushed the sujinesh/add_scale_testing_scripts branch 4 times, most recently from 59799d7 to 0c9e709 Compare

January 31, 2025 22:14

SujeethJinesh force-pushed the sujinesh/add_scale_testing_scripts branch 3 times, most recently from 68bde9d to c344cd0 Compare

February 13, 2025 01:34

SujeethJinesh changed the title ~~Add Pathways Benchmarking Recipes for Scale Testing (and fix bugs)~~ Add Pathways Recipe Support for Scale Testing (and fix bugs)

SujeethJinesh force-pushed the sujinesh/add_scale_testing_scripts branch 4 times, most recently from e5f9c2e to 7888e57 Compare

February 19, 2025 00:38

SujeethJinesh mentioned this pull request

Sync main commits with develop branch AI-Hypercomputer/xpk#379

Merged

SujeethJinesh force-pushed the sujinesh/add_scale_testing_scripts branch 3 times, most recently from ee2a336 to 72df4d8 Compare

February 19, 2025 21:39

SujeethJinesh changed the title ~~Add Pathways Recipe Support for Scale Testing (and fix bugs)~~ Add Pathways Recipe Support for Scale Testing

SujeethJinesh force-pushed the sujinesh/add_scale_testing_scripts branch from 72df4d8 to c0b018d Compare

February 19, 2025 21:59

SujeethJinesh requested a review from shauryagup

February 19, 2025 21:59

SujeethJinesh force-pushed the sujinesh/add_scale_testing_scripts branch from c0b018d to b6c7658 Compare

February 19, 2025 22:44

SujeethJinesh marked this pull request as ready for review

February 19, 2025 22:44

SujeethJinesh requested review from gobbleturk, khatwanimohit, bvandermoon, vipannalla and RissyRan as code owners

February 19, 2025 22:44

SujeethJinesh requested a review from RoshaniN

February 19, 2025 22:45

SujeethJinesh force-pushed the sujinesh/add_scale_testing_scripts branch from e070b70 to 1c1dc1e Compare

February 21, 2025 18:12

RoshaniN reviewed

View reviewed changes

Collaborator

RoshaniN left a comment

Thanks Sujeeth!

benchmarks/maxtext_xpk_runner.py Outdated Show resolved Hide resolved

benchmarks/maxtext_xpk_runner.py Outdated Show resolved Hide resolved

benchmarks/maxtext_xpk_runner.py Show resolved Hide resolved

benchmarks/recipes/pw_long_running_recipe.py Show resolved Hide resolved

benchmarks/recipes/pw_long_running_recipe.py Show resolved Hide resolved

benchmarks/recipes/pw_mcjax_benchmark_recipe.py Show resolved Hide resolved

benchmarks/recipes/pw_mcjax_benchmark_recipe.py Outdated Show resolved Hide resolved

benchmarks/xla_flags_library.py Show resolved Hide resolved

SujeethJinesh force-pushed the sujinesh/add_scale_testing_scripts branch 3 times, most recently from 965eb9a to 2279475 Compare

February 21, 2025 22:00

shauryagup reviewed

View reviewed changes

Collaborator

shauryagup left a comment

Thanks Sujeeth!

benchmarks/maxtext_trillium_model_configs.py Show resolved Hide resolved

benchmarks/maxtext_trillium_model_configs.py Outdated Show resolved Hide resolved

benchmarks/maxtext_trillium_model_configs.py Outdated Show resolved Hide resolved

benchmarks/maxtext_trillium_model_configs.py Outdated Show resolved Hide resolved

benchmarks/maxtext_trillium_model_configs.py Outdated Show resolved Hide resolved

benchmarks/maxtext_trillium_model_configs.py Outdated Show resolved Hide resolved

benchmarks/recipes/pw_long_running_recipe.py Outdated Show resolved Hide resolved

benchmarks/recipes/pw_long_running_recipe.py Outdated Show resolved Hide resolved

benchmarks/recipes/pw_mcjax_benchmark_recipe.py Outdated Show resolved Hide resolved

SujeethJinesh force-pushed the sujinesh/add_scale_testing_scripts branch 4 times, most recently from c292d23 to d2e9450 Compare

February 23, 2025 05:37

SujeethJinesh commented

View reviewed changes

Collaborator Author

SujeethJinesh left a comment

Thanks for the review!

benchmarks/maxtext_trillium_model_configs.py Show resolved Hide resolved

benchmarks/maxtext_trillium_model_configs.py Outdated Show resolved Hide resolved

benchmarks/maxtext_trillium_model_configs.py Outdated Show resolved Hide resolved

benchmarks/maxtext_trillium_model_configs.py Outdated Show resolved Hide resolved

benchmarks/recipes/pw_long_running_recipe.py Outdated Show resolved Hide resolved

SujeethJinesh force-pushed the sujinesh/add_scale_testing_scripts branch from d2e9450 to 052f650 Compare

February 25, 2025 21:43

SujeethJinesh requested review from richjames0, rni418 and gagika as code owners

February 25, 2025 21:43

SujeethJinesh force-pushed the sujinesh/add_scale_testing_scripts branch 3 times, most recently from d40bff8 to d1fdb9c Compare

February 28, 2025 22:48

shauryagup approved these changes

View reviewed changes

benchmarks/maxtext_xpk_runner.py Show resolved Hide resolved

SujeethJinesh force-pushed the sujinesh/add_scale_testing_scripts branch from d1fdb9c to 8eba24b Compare

February 28, 2025 23:04

dipannita08 approved these changes

View reviewed changes

Collaborator

dipannita08 left a comment

Looks good, thanks so much!

gobbleturk approved these changes

View reviewed changes

benchmarks/recipes/pw_long_running_recipe.py Show resolved Hide resolved

benchmarks/recipes/pw_mcjax_benchmark_recipe.py Show resolved Hide resolved

SujeethJinesh force-pushed the sujinesh/add_scale_testing_scripts branch from 8eba24b to a410df6 Compare

February 28, 2025 23:35

github-actions bot added the pull ready label

dipannita08 reviewed

View reviewed changes

benchmarks/recipes/pw_long_running_recipe.py Show resolved Hide resolved


          Add Pathways Benchmarking Recipes for Scale Testing

f164e8b

Enable metrics collection

Revert "Enable metrics collection"

This reverts commit b3d2fc0.

SujeethJinesh force-pushed the sujinesh/add_scale_testing_scripts branch from 8d4b5ea to f164e8b Compare

March 1, 2025 07:39

Collaborator Author

SujeethJinesh commented Mar 1, 2025

Closing as this PR has been merged in as part of another commit.

SujeethJinesh closed this

bvandermoon mentioned this pull request

Fix VertexAI TB Fails #1332

Merged

4 tasks

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

dipannita08 dipannita08 approved these changes

RoshaniN RoshaniN left review comments

shauryagup shauryagup approved these changes

gobbleturk gobbleturk approved these changes

khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit is a code owner

bvandermoon Awaiting requested review from bvandermoon bvandermoon is a code owner

vipannalla Awaiting requested review from vipannalla vipannalla is a code owner

RissyRan Awaiting requested review from RissyRan RissyRan is a code owner

richjames0 Awaiting requested review from richjames0 richjames0 is a code owner

rni418 Awaiting requested review from rni418

gagika Awaiting requested review from gagika gagika is a code owner

Labels