Skip to content

Add SPMD config option to specify zero cost method for gather/scatter. - #14862

Closed
patrick-toulme wants to merge 1 commit into
openxla:mainfrom
patrick-toulme:gather_scatter_config
Closed

patrick-toulme wants to merge 1 commit into
openxla:mainfrom
patrick-toulme:gather_scatter_config

Conversation

@patrick-toulme

@patrick-toulme patrick-toulme commented Jul 12, 2024 •

Copy link
Copy Markdown
Contributor

Issue #13304

In SPMD handling of gather/scatter the partition strategy is hardcoded to IndexParallel strategy. This is not optimal for all topology. This PR makes this option an SPMD config, but defaults to IndexParallel to maintain existing behavior.

Clang-format also fixed some formatting. Tests were added and all tests pass.

@patrick-toulme
patrick-toulme force-pushed the gather_scatter_config branch from b633c22 to 7f83c21 Compare July 12, 2024 18:11
@patrick-toulme

patrick-toulme commented Jul 12, 2024 •

Copy link
Copy Markdown
Contributor Author

@frgossen please review. Thank you.

@NaiyerRizz NaiyerRizz self-assigned this Jul 13, 2024
@NaiyerRizz
NaiyerRizz requested a review from frgossen July 13, 2024 04:27
@thomasjoerg

Copy link
Copy Markdown
Member

@frgossen friendly reminder!

@frgossen frgossen left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Clang-format also fixed some formatting.

These PRs will be auto-formatted on import so those changes might disappear again. Generally, it will be easier to review if you don't change formatting.

Thanks for making this configurable!

copybara-service Bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jul 15, 2024
…her/scatter.

Imported from GitHub PR openxla/xla#14862

Issue #13304

In SPMD handling of gather/scatter the partition strategy is hardcoded to IndexParallel strategy. This is not optimal for all topology. This PR makes this option an SPMD config, but defaults to IndexParallel to maintain existing behavior.

Clang-format also fixed some formatting. Tests were added and all tests pass.
Copybara import of the project:

--
7f83c21573f24cd4e314b13ce2e349dd6194b451 by ptoulme-aws <[email protected]>:

Add SPMD config option to specify zero cost method for gather/scatter.

Merging this change closes #14862

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14862 from ptoulme-aws:gather_scatter_config 7f83c21573f24cd4e314b13ce2e349dd6194b451
PiperOrigin-RevId: 652473292
copybara-service Bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jul 15, 2024
…her/scatter.

Imported from GitHub PR openxla/xla#14862

Issue #13304

In SPMD handling of gather/scatter the partition strategy is hardcoded to IndexParallel strategy. This is not optimal for all topology. This PR makes this option an SPMD config, but defaults to IndexParallel to maintain existing behavior.

Clang-format also fixed some formatting. Tests were added and all tests pass.
Copybara import of the project:

--
7f83c21573f24cd4e314b13ce2e349dd6194b451 by ptoulme-aws <[email protected]>:

Add SPMD config option to specify zero cost method for gather/scatter.

Merging this change closes #14862

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14862 from ptoulme-aws:gather_scatter_config 7f83c21573f24cd4e314b13ce2e349dd6194b451
PiperOrigin-RevId: 652473292
copybara-service Bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jul 15, 2024
…her/scatter.

Imported from GitHub PR openxla/xla#14862

Issue #13304

In SPMD handling of gather/scatter the partition strategy is hardcoded to IndexParallel strategy. This is not optimal for all topology. This PR makes this option an SPMD config, but defaults to IndexParallel to maintain existing behavior.

Clang-format also fixed some formatting. Tests were added and all tests pass.
Copybara import of the project:

--
7f83c21573f24cd4e314b13ce2e349dd6194b451 by ptoulme-aws <[email protected]>:

Add SPMD config option to specify zero cost method for gather/scatter.

Merging this change closes #14862

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14862 from ptoulme-aws:gather_scatter_config 7f83c21573f24cd4e314b13ce2e349dd6194b451
PiperOrigin-RevId: 652473292
copybara-service Bot pushed a commit that referenced this pull request Jul 15, 2024
…her/scatter.

Imported from GitHub PR #14862

Issue #13304

In SPMD handling of gather/scatter the partition strategy is hardcoded to IndexParallel strategy. This is not optimal for all topology. This PR makes this option an SPMD config, but defaults to IndexParallel to maintain existing behavior.

Clang-format also fixed some formatting. Tests were added and all tests pass.
Copybara import of the project:

--
7f83c21 by ptoulme-aws <[email protected]>:

Add SPMD config option to specify zero cost method for gather/scatter.

Merging this change closes #14862

FUTURE_COPYBARA_INTEGRATE_REVIEW=#14862 from ptoulme-aws:gather_scatter_config 7f83c21
PiperOrigin-RevId: 652514718
copybara-service Bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jul 15, 2024
…her/scatter.

Imported from GitHub PR openxla/xla#14862

Issue #13304

In SPMD handling of gather/scatter the partition strategy is hardcoded to IndexParallel strategy. This is not optimal for all topology. This PR makes this option an SPMD config, but defaults to IndexParallel to maintain existing behavior.

Clang-format also fixed some formatting. Tests were added and all tests pass.
Copybara import of the project:

--
7f83c21573f24cd4e314b13ce2e349dd6194b451 by ptoulme-aws <[email protected]>:

Add SPMD config option to specify zero cost method for gather/scatter.

Merging this change closes #14862

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14862 from ptoulme-aws:gather_scatter_config 7f83c21573f24cd4e314b13ce2e349dd6194b451
PiperOrigin-RevId: 652514718
@patrick-toulme

Copy link
Copy Markdown
Contributor Author

Clang-format also fixed some formatting.

These PRs will be auto-formatted on import so those changes might disappear again. Generally, it will be easier to review if you don't change formatting.

Thanks for making this configurable!

Okay, I will not format any future changes. Thank you for reviewing.

I see my PR is not merging due to ROCm and Windows Bazel build failures in copy-bara import to Tensorflow. I cannot view the ROCm failure. The Windows failure is on this target - tensorflow/tools/ci_build/builds/gen_win_out.exe. How do I debug this unless these builds are broken on TOT right now?

copybara-service Bot pushed a commit that referenced this pull request Jul 16, 2024
…her/scatter.

Imported from GitHub PR #14862

Issue #13304

In SPMD handling of gather/scatter the partition strategy is hardcoded to IndexParallel strategy. This is not optimal for all topology. This PR makes this option an SPMD config, but defaults to IndexParallel to maintain existing behavior.

Clang-format also fixed some formatting. Tests were added and all tests pass.
Copybara import of the project:

--
7f83c21 by ptoulme-aws <[email protected]>:

Add SPMD config option to specify zero cost method for gather/scatter.

Merging this change closes #14862

FUTURE_COPYBARA_INTEGRATE_REVIEW=#14862 from ptoulme-aws:gather_scatter_config 7f83c21
PiperOrigin-RevId: 652514718
copybara-service Bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jul 16, 2024
…her/scatter.

Imported from GitHub PR openxla/xla#14862

Issue #13304

In SPMD handling of gather/scatter the partition strategy is hardcoded to IndexParallel strategy. This is not optimal for all topology. This PR makes this option an SPMD config, but defaults to IndexParallel to maintain existing behavior.

Clang-format also fixed some formatting. Tests were added and all tests pass.
Copybara import of the project:

--
7f83c21573f24cd4e314b13ce2e349dd6194b451 by ptoulme-aws <[email protected]>:

Add SPMD config option to specify zero cost method for gather/scatter.

Merging this change closes #14862

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14862 from ptoulme-aws:gather_scatter_config 7f83c21573f24cd4e314b13ce2e349dd6194b451
PiperOrigin-RevId: 652514718
copybara-service Bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jul 16, 2024
…her/scatter.

Imported from GitHub PR openxla/xla#14862

Issue #13304

In SPMD handling of gather/scatter the partition strategy is hardcoded to IndexParallel strategy. This is not optimal for all topology. This PR makes this option an SPMD config, but defaults to IndexParallel to maintain existing behavior.

Clang-format also fixed some formatting. Tests were added and all tests pass.
Copybara import of the project:

--
7f83c21573f24cd4e314b13ce2e349dd6194b451 by ptoulme-aws <[email protected]>:

Add SPMD config option to specify zero cost method for gather/scatter.

Merging this change closes #14862

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14862 from ptoulme-aws:gather_scatter_config 7f83c21573f24cd4e314b13ce2e349dd6194b451
PiperOrigin-RevId: 652473292
copybara-service Bot pushed a commit that referenced this pull request Jul 16, 2024
FUTURE_COPYBARA_INTEGRATE_REVIEW=#14862 from ptoulme-aws:gather_scatter_config 7f83c21
PiperOrigin-RevId: 652312882
copybara-service Bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jul 16, 2024
FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14862 from ptoulme-aws:gather_scatter_config 7f83c21573f24cd4e314b13ce2e349dd6194b451
PiperOrigin-RevId: 652312882
copybara-service Bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jul 16, 2024
FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14862 from ptoulme-aws:gather_scatter_config 7f83c21573f24cd4e314b13ce2e349dd6194b451
PiperOrigin-RevId: 652134750
copybara-service Bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jul 16, 2024
FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14862 from ptoulme-aws:gather_scatter_config 7f83c21573f24cd4e314b13ce2e349dd6194b451
PiperOrigin-RevId: 652312888
copybara-service Bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jul 16, 2024
FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14862 from ptoulme-aws:gather_scatter_config 7f83c21573f24cd4e314b13ce2e349dd6194b451
PiperOrigin-RevId: 652312889
copybara-service Bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jul 16, 2024
FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14862 from ptoulme-aws:gather_scatter_config 7f83c21573f24cd4e314b13ce2e349dd6194b451
PiperOrigin-RevId: 652135040
copybara-service Bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jul 16, 2024
FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14862 from ptoulme-aws:gather_scatter_config 7f83c21573f24cd4e314b13ce2e349dd6194b451
PiperOrigin-RevId: 650565685
copybara-service Bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jul 16, 2024
FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14862 from ptoulme-aws:gather_scatter_config 7f83c21573f24cd4e314b13ce2e349dd6194b451
PiperOrigin-RevId: 652312897
copybara-service Bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jul 16, 2024
FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14862 from ptoulme-aws:gather_scatter_config 7f83c21573f24cd4e314b13ce2e349dd6194b451
PiperOrigin-RevId: 651940791
copybara-service Bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jul 16, 2024
FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14862 from ptoulme-aws:gather_scatter_config 7f83c21573f24cd4e314b13ce2e349dd6194b451
PiperOrigin-RevId: 651940797
copybara-service Bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jul 16, 2024
FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#14862 from ptoulme-aws:gather_scatter_config 7f83c21573f24cd4e314b13ce2e349dd6194b451
PiperOrigin-RevId: 652312894
copybara-service Bot pushed a commit to tensorflow/tensorflow that referenced this pull request Jul 16, 2024
…her/scatter.

Imported from GitHub PR openxla/xla#14862

Issue #13304

In SPMD handling of gather/scatter the partition strategy is hardcoded to IndexParallel strategy. This is not optimal for all topology. This PR makes this option an SPMD config, but defaults to IndexParallel to maintain existing behavior.

Clang-format also fixed some formatting. Tests were added and all tests pass.
Copybara import of the project:

--
7f83c21573f24cd4e314b13ce2e349dd6194b451 by ptoulme-aws <[email protected]>:

Add SPMD config option to specify zero cost method for gather/scatter.

Merging this change closes #14862

PiperOrigin-RevId: 652736743
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants