Align Int4Tensor implementation details with the design of Float8Tensor

[jerryzh168]

Stacked PRs:

• Support optional_tensor_names in TorchAOBaseTensor #2710
• ->Align Int4Tensor implementation details with the design of Float8Tensor #2687
• Add all fbgemm kernel Tensors into Int4WeightOnlyConfig and Float8DynamicActivationInt4WeightConfig #2474
• Check numerical equivalence / closeness between different kernel preferences #2651

---

Align Int4Tensor implementation details with the design of Float8Tensor

Summary:
Int4Tensor is the non-preshuffled version of int4 quantized Tensor, data is [N, K/2], scale/zero_point has shape: [K/group_size, N]

Multiple fixes for Int4Tensor to align with the design of Float8Tensor (only calling fbgemm ops)

• Added VERSION 2 for Int4WeightOnlyConfig
• Migrated op implementation and tests from Add support for resharding and int4 preshuffle kernel #2387

Test Plan:
python test/quantization/quantize_/workflows/int4/test_int4_tensor.py

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⏳ No Failures, 1 Pending

As of commit f9695e4 with merge base d2e791b ():
💚 Looks good so far! There are no failures yet. 💚

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

[vkuzo]

why rename this? from_float seems cleaner, maybe from_high_precision might be even more clear. Same feedback for Float8Tensor.

[jerryzh168]

this is to align with Float8Tensor, sure, can rename this to from_high_precision for both

[jerryzh168]

how about from_float_hp? this is a bot more precise than high_precision I think

[vkuzo]

is it expected that the tensors are not contiguous? if not, can we assert for this instead of calling contiguous?

[jerryzh168]

the non-contiguous comes from the reshape ops like transpose, view I think, but the kernel will need these to be contiguous, I can try changing these to assert and do the contiguous operation in user side to see if it works

[vkuzo]

I would have expected the weights to be stored in a format aligned with what the kernel needs, without any need for just-in-time layout transforms. Does this match how the current code works?

[jerryzh168]

normally it is, but the weights also goes through some transformations like the ones we listed in test_moe_weight_reshape_ops which makes weight / scale etc. non-contiguous I think, but I can try to do call contiguous in user code, that might be cleaner I think

[vkuzo]

this is really confusing, could you share the result of print(model) after this function has been applied?

if it's going to print model with parameters wrapped in Int4Tensor, can we just wrap the parameters directly without all of these layers of abstraction?

[vkuzo]

if this is working around the fact that quantize_ needs to work on modules, IMO we should change quantize_ to handle this instead of working around? seems important for MoEs.

[jerryzh168]

yeah the parameters are wrapped in Int4Tensor, this is just applying quantization to each of the moe weights: w1, w2 and w3

I can inline these for now. can follow up with how to have an API for weights + configs separately

[vkuzo]

maybe call it something like FeedForwardWithExperts? Experts is ambiguous

[jerryzh168]

this is adapted from https://github.com/meta-llama/llama-models/blob/a9c89c471f793423afd4cc3ca8671d6e56fe64cb/models/llama4/moe.py#L22, how about renaming to LLama4Experts to make it more specific