Add layout option to woq int4 api

[Diogo-V]

Summary

This PR updates the int4_weight_only API by introducing a layout_type parameter. This will allow users to quantize models to int4 while being able to choose between available layouts.

This change also breaks backwards compatibility for all users that explicitly define the inner_k_tiles parameter:

# old
def int4_weight_only(group_size=128, inner_k_tiles=8):
...

quantize_(my_model, int4_weight_only(inner_k_tiles=8))

# new
def int4_weight_only(group_size=128, layout_type=TensorCoreTiledLayoutType(inner_k_tiles=8)):
...

quantize_(my_model, int4_weight_only(layout_type=TensorCoreTiledLayoutType(inner_k_tiles=8)))

BC Breaking notes for Release statement

• inner_k_tiles was deprecated in favor of layout_type, enabling users to select from various layout options #670

# for torchao 0.5
from torchao.quantization import quantize, int4_weight_only
quantize_(my_model, int4_weight_only(layout_type=TensorCoreTiledLayoutType(inner_k_tiles=8)))

# for torchao 0.4
from torchao.quantization import quantize_, int4_weight_only
quantize_(my_model, int4_weight_only(inner_k_tiles=8))

Tasks

• Update the function's interface
• Update the documentation
• Update the tests

PRs pending on this one

• Add sparse marlin AQT layout #621

Questions for Reviewers

• I am not familiar with how releases and backwards compatibility is maintained in pytorch and since this is a breaking change, how should we handle delivering it? Is there a specific flag/if statement that I should add somewhere?

Let me know if there is anything that I need to update. I would be happy to perform any necessary changes.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 500a456 with merge base 88a263a ():
💚 Looks good so far! There are no failures yet. 💚

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

[jerryzh168]

thanks, if there is no circular dep you can remove the import from other functions as well, e.g. int8_weight_only

[Diogo-V]

done

[jerryzh168]

thanks, although this is bc-breaking, but I don't think there are many callsites changing inner_k_tiles so it probably won't affect many people I think

[jerryzh168]

could you write a bc-breaking notes for this? see format in BC-breaking section of https://github.com/pytorch/ao/releases, something like:

# for torchao 0.5
...

# for torchao 0.4
...

[Diogo-V]

on it! where should I place the bc-breaking notes? as a comment on this PR or someplace else?

[Diogo-V]

• inner_k_tiles was deprecated in favor of layout_type, enabling users to select from various layout options #670

# for torchao 0.5
from torchao.quantization import quantize_, int4_weight_only
quantize_(my_model, int4_weight_only(inner_k_tiles=8))

# for torchao 0.4
from torchao.quantization import quantize, int4_weight_only
quantize_(my_model, int4_weight_only(layout_type=TensorCoreTiledLayoutType(inner_k_tiles=8)))

[jerryzh168]

on it! where should I place the bc-breaking notes? as a comment on this PR or someplace else?

you can put this in PR summary

[jerryzh168]

looks good, thanks!

[Diogo-V]

done ✅
feel free to merge it or let me know if there is anything else that I might need to change

[msaroufim]

@jerryzh168 do we need to communicate this change to anyone?

[jerryzh168]

@jerryzh168 do we need to communicate this change to anyone?

seems fine I think, I didn't see people using inner_k_tiles yet, but we can probably update huggingface when 0.5 is released