Implement operator for micro-pipelined all-gather -> _scaled_mm #129289

yifuwang · 2024-06-22T00:09:57Z

Stack from ghstack (oldest at bottom):

This PR implements torch.ops.symm_mem.fused_all_gather_scaled_matmul. It's similar to torch.ops.symm_mem.fused_all_gather_matmul, except that it takes scales and calls _scaled_mm.

Profiling Trace vs. Baseline (FB internal only)

cc @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @penguinwu @tianyu-l @yf225 @chauhang

[ghstack-poisoned]

pytorch-bot · 2024-06-22T00:10:00Z

🔗 Helpful Links

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

📄 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.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 4a81ea1 with merge base a21d436 ():

UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:

trunk / win-vs2019-cpu-py3 / test (default, 2, 3, windows.4xlarge.nonephemeral, unstable) (gh)
test_decomp 2/12 failed!

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

ghstack-source-id: 268c246f430ed35d4186822e9a55320ca174f1c5 Pull Request resolved: #129289

weifengpy · 2024-06-24T17:39:25Z

fyi @drisspg on matmuls in async TP

drisspg

Very cool it can all be done in python

…ed_mm" This PR implements `torch.ops.symm_mem.fused_all_gather_scaled_matmul`. It's similar to `torch.ops.symm_mem.fused_all_gather_matmul`, except that it takes scales and calls ` _scaled_mm`. [Profiling Trace vs. Baseline](https://interncache-all.fbcdn.net/manifold/perfetto-artifacts/tree/ui/index.html#!/?url=https://interncache-all.fbcdn.net/manifold/perfetto_internal_traces/tree/shared_trace/yifu_tmp0gmg1f2_) (FB internal only) cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k [ghstack-poisoned]

ghstack-source-id: 08dec6ca9027f97e6252fb863442b21a45aa3552 Pull Request resolved: #129289

…ed_mm" This PR implements `torch.ops.symm_mem.fused_all_gather_scaled_matmul`. It's similar to `torch.ops.symm_mem.fused_all_gather_matmul`, except that it takes scales and calls ` _scaled_mm`. [Profiling Trace vs. Baseline](https://interncache-all.fbcdn.net/manifold/perfetto-artifacts/tree/ui/index.html#!/?url=https://interncache-all.fbcdn.net/manifold/perfetto_internal_traces/tree/shared_trace/yifu_tmp0gmg1f2_) (FB internal only) cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k [ghstack-poisoned]

ghstack-source-id: d7e7bd2c6a386f0574e052130daa38ec0b55dc91 Pull Request resolved: #129289

yifuwang · 2024-06-28T20:45:38Z

@pytorchbot merge

pytorchmergebot · 2024-06-28T20:47:18Z

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 · 2024-06-29T02:46:07Z

The merge job was canceled or timed out. This most often happen if two merge requests were issued for the same PR, or if merge job was waiting for more than 6 hours for tests to finish. In later case, please do not hesitate to reissue the merge command
For more information see pytorch-bot wiki.

yifuwang · 2024-07-01T18:29:04Z

@pytorchbot merge

pytorchmergebot · 2024-07-01T18:30:41Z

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 · 2024-07-02T00:29:32Z

The merge job was canceled or timed out. This most often happen if two merge requests were issued for the same PR, or if merge job was waiting for more than 6 hours for tests to finish. In later case, please do not hesitate to reissue the merge command
For more information see pytorch-bot wiki.

…ed_mm" This PR implements `torch.ops.symm_mem.fused_all_gather_scaled_matmul`. It's similar to `torch.ops.symm_mem.fused_all_gather_matmul`, except that it takes scales and calls ` _scaled_mm`. [Profiling Trace vs. Baseline](https://interncache-all.fbcdn.net/manifold/perfetto-artifacts/tree/ui/index.html#!/?url=https://interncache-all.fbcdn.net/manifold/perfetto_internal_traces/tree/shared_trace/yifu_tmp0gmg1f2_) (FB internal only) cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k [ghstack-poisoned]

[ghstack-poisoned]

ghstack-source-id: 189cbb29deeaadd1982a287147ff561926ae9268 Pull Request resolved: pytorch#129289

yifuwang · 2024-07-15T21:46:36Z

@pytorchbot merge

pytorchmergebot · 2024-07-15T21:48:22Z

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

…rch#129289) This PR implements `torch.ops.symm_mem.fused_all_gather_scaled_matmul`. It's similar to `torch.ops.symm_mem.fused_all_gather_matmul`, except that it takes scales and calls ` _scaled_mm`. [Profiling Trace vs. Baseline](https://interncache-all.fbcdn.net/manifold/perfetto-artifacts/tree/ui/index.html#!/?url=https://interncache-all.fbcdn.net/manifold/perfetto_internal_traces/tree/shared_trace/yifu_tmp0gmg1f2_) (FB internal only) Co-authored-by: Will Feng <[email protected]> Pull Request resolved: pytorch#129289 Approved by: https://github.com/Chillee, https://github.com/weifengpy, https://github.com/drisspg

Implement operator for micro-pipelined all-gather -> _scaled_mm

01a987f

[ghstack-poisoned]

yifuwang mentioned this pull request Jun 21, 2024

Remove ProcessGroupCudaP2P and change async-TP to use SymmetricMemory #128762

Closed

pytorch-bot bot added the oncall: distributed Add this issue/PR to distributed oncall triage queue label Jun 22, 2024

yifuwang added a commit that referenced this pull request Jun 22, 2024

Implement operator for micro-pipelined all-gather -> _scaled_mm

4103331

ghstack-source-id: 268c246f430ed35d4186822e9a55320ca174f1c5 Pull Request resolved: #129289

yifuwang requested review from wanchaol, Chillee and weifengpy June 22, 2024 00:10

Chillee approved these changes Jun 22, 2024

View reviewed changes

weifengpy approved these changes Jun 24, 2024

View reviewed changes

drisspg approved these changes Jun 24, 2024

View reviewed changes

yifuwang added a commit that referenced this pull request Jun 25, 2024

Implement operator for micro-pipelined all-gather -> _scaled_mm

ec42c47

ghstack-source-id: 08dec6ca9027f97e6252fb863442b21a45aa3552 Pull Request resolved: #129289

yifuwang added a commit that referenced this pull request Jun 28, 2024

Implement operator for micro-pipelined all-gather -> _scaled_mm

6cd4bbb

ghstack-source-id: d7e7bd2c6a386f0574e052130daa38ec0b55dc91 Pull Request resolved: #129289

yifuwang added the topic: not user facing topic category label Jun 28, 2024

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Jun 28, 2024

pytorchmergebot added the merging label Jun 28, 2024

yifuwang mentioned this pull request Jul 2, 2024

Fix and improve raise_comms and sink_waits #129980

Closed

Update

4a81ea1

[ghstack-poisoned]

yf225 mentioned this pull request Jul 8, 2024

fix reorder pass algo after rebase #130276

Closed

OnlyFor pushed a commit to OnlyFor/pytorch that referenced this pull request Jul 9, 2024

Implement operator for micro-pipelined all-gather -> _scaled_mm

dba15ee

ghstack-source-id: 189cbb29deeaadd1982a287147ff561926ae9268 Pull Request resolved: pytorch#129289

pytorchmergebot closed this in db3a641 Jul 15, 2024

pytorchmergebot added Merged and removed merging labels Jul 15, 2024

henrylhtsang mentioned this pull request Jul 17, 2024

[aoti] Unskip some aot inductor tests #130973

Closed

github-actions bot deleted the gh/yifuwang/95/head branch August 15, 2024 01:54

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement operator for micro-pipelined all-gather -> _scaled_mm #129289

Implement operator for micro-pipelined all-gather -> _scaled_mm #129289

yifuwang commented Jun 22, 2024 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Jun 22, 2024 •

edited

Loading

weifengpy commented Jun 24, 2024

drisspg left a comment

yifuwang commented Jun 28, 2024

pytorchmergebot commented Jun 28, 2024

pytorchmergebot commented Jun 29, 2024

yifuwang commented Jul 1, 2024

pytorchmergebot commented Jul 1, 2024

pytorchmergebot commented Jul 2, 2024

yifuwang commented Jul 15, 2024

pytorchmergebot commented Jul 15, 2024

Implement operator for micro-pipelined all-gather -> _scaled_mm #129289

Implement operator for micro-pipelined all-gather -> _scaled_mm #129289

Conversation

yifuwang commented Jun 22, 2024 • edited by pytorch-bot bot Loading

pytorch-bot bot commented Jun 22, 2024 • edited Loading

🔗 Helpful Links

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

✅ You can merge normally! (1 Unrelated Failure)

weifengpy commented Jun 24, 2024

drisspg left a comment

Choose a reason for hiding this comment

yifuwang commented Jun 28, 2024

pytorchmergebot commented Jun 28, 2024

Merge started

pytorchmergebot commented Jun 29, 2024

yifuwang commented Jul 1, 2024

pytorchmergebot commented Jul 1, 2024

Merge started

pytorchmergebot commented Jul 2, 2024

yifuwang commented Jul 15, 2024

pytorchmergebot commented Jul 15, 2024

Merge started

yifuwang commented Jun 22, 2024 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Jun 22, 2024 •

edited

Loading