Skip to content

nvfp4 gptq for bmm#4327

Merged
vkuzo merged 58 commits intomainfrom
gh/vkuzo/260/head
Apr 27, 2026
Merged

nvfp4 gptq for bmm#4327
vkuzo merged 58 commits intomainfrom
gh/vkuzo/260/head

Conversation

@vkuzo
Copy link
Copy Markdown
Contributor

@vkuzo vkuzo commented Apr 24, 2026

Summary:

Extend GPTQ coverage for bmm, formulating the bmm as a 3d case of mm.
This involves:

  1. refactoring the 2d code to make it easily extendable to 3d
  2. the existing bmm logic was numerically incorrect (used a single
    hessian), modify it to instead use E K by K Hessians for an E, N, K
    input shape, route to the 2D hessian logic E times. This is slow
    but we can optimize later.

We test numerical correctness by bitwise matching E 2d hessian
calculations to the 3D one.

Test Plan:

torchao/prototype/gptq/gptq_nvfp4_llama3_2_1b_nonsequential_wikitext.sh
// gptq accuracy unchanged

vkuzo added 30 commits April 20, 2026 20:52
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
vkuzo added 10 commits April 23, 2026 18:44
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
@vkuzo
Copy link
Copy Markdown
Contributor Author

vkuzo commented Apr 24, 2026

@vkuzo vkuzo requested a review from jerryzh168 as a code owner April 24, 2026 12:52
@pytorch-bot
Copy link
Copy Markdown

pytorch-bot Bot commented Apr 24, 2026

🔗 Helpful Links

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

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:

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

@meta-cla meta-cla 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 Apr 24, 2026
@vkuzo vkuzo added the module: not user facing Use this tag if you don't want this PR to show up in release notes label Apr 24, 2026
assert observer.total_batches == 0

@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA available")
def test_observer_tensor_attributes(self):
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

this test is not that useful, deleting instead of updating with new contents of the observer tensor

vkuzo added 2 commits April 24, 2026 14:04
[ghstack-poisoned]
[ghstack-poisoned]
return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined]

def __init__(self, hp_data: torch.Tensor, total_batches: int, hessian=None):
def __init__(self, hp_data: torch.Tensor, total_batches, hessian=None):
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

can you add some docs for total_batches?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

sure, let me do that in a future PR (planning today), as we also need to change the definition of this for grouped_mm to sample each token equally instead of sampling each batch equally

[ghstack-poisoned]
@vkuzo vkuzo requested a review from danielvegamyhre as a code owner April 27, 2026 12:07
@vkuzo vkuzo changed the base branch from gh/vkuzo/259/head to main April 27, 2026 12:07
@vkuzo vkuzo merged commit f9a5f28 into main Apr 27, 2026
38 of 66 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. module: not user facing Use this tag if you don't want this PR to show up in release notes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants