[ROCm][SymmetricMemory] Performance improvements for two-shot allreduce

[pragupta]

The biggest bottleneck that we found with two-shot allreduce was that the compiler was serializing all the load operations for some reason. To avoid these load delays, we've added de-serialization of loads. Along with this improvement, we also found that on AMD GPUs a different block and thread size gives a nice performance boost. Here are the bandwidth numbers I am getting with this PR:

The rows that are green are the tensor sizes that we are interested in because two-shot is only used for bigger sizes (one-shot is used for smaller sizes). As we can see, our baseline numbers wrt to fbgemm numbers were consistently underperforming. However, with this deserialize change, most of the tensor sizes have a performance boost (positive %) for the green tensors. There's one tensor with negative performance, but that's within error margin.

co-authored by: @amd-hhashemi
pytorch/FBGEMM#4072

cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @dllehr-amd @jataylo @hongxiayang @naromero77amd

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/156746

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 6 New Failures, 1 Pending

As of commit e4c42ba with merge base aa280ea ():

NEW FAILURES - The following jobs have failed:

• pull / linux-jammy-py3.13-clang12 / test (dynamo_wrapped, 1, 3, linux.2xlarge) (gh)
  Found the unimplemented_v2 or unimplemented_v2_with_warning calls below that don't match the registry in graph_break_registry.json.
• pull / linux-jammy-py3.13-clang12 / test (dynamo_wrapped, 2, 3, linux.2xlarge) (gh)
  Found the unimplemented_v2 or unimplemented_v2_with_warning calls below that don't match the registry in graph_break_registry.json.
• pull / linux-jammy-py3.13-clang12 / test (dynamo_wrapped, 3, 3, linux.2xlarge) (gh)
  Found the unimplemented_v2 or unimplemented_v2_with_warning calls below that don't match the registry in graph_break_registry.json.
• pull / linux-jammy-py3.9-clang12 / test (dynamo_wrapped, 1, 3, linux.2xlarge) (gh)
  Found the unimplemented_v2 or unimplemented_v2_with_warning calls below that don't match the registry in graph_break_registry.json.
• pull / linux-jammy-py3.9-clang12 / test (dynamo_wrapped, 2, 3, linux.2xlarge) (gh)
  Found the unimplemented_v2 or unimplemented_v2_with_warning calls below that don't match the registry in graph_break_registry.json.
• pull / linux-jammy-py3.9-clang12 / test (dynamo_wrapped, 3, 3, linux.2xlarge) (gh)
  Found the unimplemented_v2 or unimplemented_v2_with_warning calls below that don't match the registry in graph_break_registry.json.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[pytorch-bot]

To add the ciflow label ciflow/trunk please first approve the workflows that are awaiting approval (scroll to the bottom of this page).

This helps ensure we don't trigger CI on this PR until it is actually authorized to do so. Please ping one of the reviewers if you do not have access to approve and run workflows.

[jeffdaily]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[pytorchmergebot]

Merge failed

Reason: 6 mandatory check(s) failed. The first few are:

• pull / linux-jammy-py3.13-clang12 / test (dynamo_wrapped, 1, 3, linux.2xlarge)
• pull / linux-jammy-py3.13-clang12 / test (dynamo_wrapped, 3, 3, linux.2xlarge)
• pull / linux-jammy-py3.13-clang12 / test (dynamo_wrapped, 2, 3, linux.2xlarge)
• pull / linux-jammy-py3.9-clang12 / test (dynamo_wrapped, 3, 3, linux.2xlarge)
• pull / linux-jammy-py3.9-clang12 / test (dynamo_wrapped, 2, 3, linux.2xlarge)

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

[jeffdaily]

@pytorchbot merge -i "rocm-only change, all rocm CI is passing"

[pytorch-bot]

❌ 🤖 pytorchbot command failed:

@pytorchbot: error: unrecognized arguments: rocm-only change, all rocm CI is passing

usage: @pytorchbot [-h] {merge,revert,rebase,label,drci,cherry-pick} ...

Try @pytorchbot --help for more info.

[jeffdaily]

@pytorchbot merge -f "rocm-only change, all rocm CI is passing"

[pytorchmergebot]

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here