8000 [NF4] Support nf4 tensor shard and gather by mori360 · Pull Request #2449 · pytorch/ao · GitHub
[go: up one dir, main page]
More Web Proxy on the site http://driver.im/
Skip to content

[NF4] Support nf4 tensor shard and gather #2449

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
Draft
wants to merge 5 commits into
base: main
Choose a base branch
from
Draft

[NF4] Support nf4 tensor shard and gather #2449

wants to merge 5 commits into from
+119 −12

Conversation

mori360
Copy link
@mori360 mori360 commented Jun 26, 2025

Add nf4_all_gather_into_tensor and scatter_nf4tensor to enable dispatch of scatter and all_gather_into_tensor
Add unit test to show that nf4 tensor keeps the same after distribute and gather.

@pytorch-bot PyTorch Bot
Copy link
pytorch-bot bot commented Jun 26, 2025 •
edited
Loading

🔗 Helpful Links

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

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

❌ 1 New Failure

As of commit f781ddb with merge base 8b57afe (image):

NEW FAILURE - The following job has failed:

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

@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 26, 2025
@mori360 mori360 marked this pull request as ready for review June 26, 2025 19:30
@mori360 mori360 added the topic: for developers Use this tag if this PR is mainly developer facing label Jun 26, 2025
@mori360 mori360 marked this pull request as draft June 26, 2025 21:06
@msaroufim
Copy link
Member
msaroufim commented Jun 26, 2025 •
edited
Loading

Presumably we should observe a comms time reduction? would be nice to see some profiles

@mori360
Copy link
Author
mori360 commented Jun 26, 2025

Presumably we should observe a comms time reduction? would be nice to see some profiles

Are there any baseline I could compare with?

@drisspg drisspg requested a review from weifengpy June 26, 2025 23:28
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@andrewor14 andrewor14 Awaiting requested review from andrewor14

@msaroufim msaroufim Awaiting requested review from msaroufim

@drisspg drisspg Awaiting requested review from drisspg

@weifengpy weifengpy Awaiting requested review from weifengpy

Assignees
No one assigned
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. topic: for developers Use this tag if this PR is mainly developer facing
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@mori360 @msaroufim @facebook-github-bot
0