Skip to content
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.

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

MX: move block_size and elem_dtype into MXLinearConfig #1689

Merged
merged 11 commits into from
Feb 14, 2025
Merged

Conversation

vkuzo
Copy link
Contributor

@vkuzo vkuzo commented Feb 10, 2025

Summary:

Moves block_size and elem_dtype into MXLinearConfig and updates all callsites.

Before

elem_dtype = torch.float8_e4m3fn
swap_linear_with_mx_linear(m, elem_dtype, block_size=32)

After

config = MXLinearConfig(elem_dtype=torch.float8_e4m3fn, block_size=32)
swap_linear_with_mx_linear(m, config=config)

Test Plan: CI

Reviewers:

Subscribers:

Tasks:

Tags:

vkuzo added 8 commits January 29, 2025 20:32
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
@vkuzo
Copy link
Contributor Author

vkuzo commented Feb 10, 2025

Copy link

pytorch-bot bot commented Feb 10, 2025

🔗 Helpful Links

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

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

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 Feb 10, 2025
vkuzo added a commit that referenced this pull request Feb 10, 2025
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 92913315b65e8b19a55bab3a51c54ed56c1bb4f4
ghstack-comment-id: 2649054194
Pull Request resolved: #1689
@vkuzo vkuzo added the topic: improvement Use this tag if this PR is an improvement (doesn't fit into any of the other categories) label Feb 10, 2025
# TODO(future PR): refactor to make this cleaner
elem_dtype_weight_override: Optional[Any] = None
elem_dtype_grad_output_override: Optional[Any] = None

# If True, uses a custom triton kernel for fp4 dequantize
use_fp4_custom_triton_dequant_kernel: bool = False
Copy link
Contributor

Choose a reason for hiding this comment

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

you think that we will want to keep this public?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

unlikely, but IMO we can punt that until later

[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
@vkuzo vkuzo changed the base branch from gh/vkuzo/25/head to main February 14, 2025 23:46
@vkuzo vkuzo merged commit 40d01cd into main Feb 14, 2025
32 of 40 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. topic: improvement Use this tag if this PR is an improvement (doesn't fit into any of the other categories)
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants