Add semi-structured sparsity to hf eval

[jcaip]

Summary:

This PR adds in huggingface eval for semi-structured sparsity with the nm-testing/SparseLlama-3-8B-pruned_50.2of4 checkpoint published by neuralmagic. We can accelerate this checkpoint with to_sparse_semi_structured with minimal accuracy loss.

Also adds in an example notebook that outlines how to do text generation with huggingface's model and shows a 80 tok/s -> 87 tok/s speedup.

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 fc82248 with merge base 5f35645 ():
💚 Looks good so far! There are no failures yet. 💚

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

[msaroufim]

how come?

[jcaip]

mkn	dense_time	sparse_time	dense_memory	sparse_memory	speedup
(1, 4096, 4096)	0.026794	0.048279	0.050628	0.136668	0.553
(1, 4096, 1024)	0.014931	0.048054	0.025443	0.046953	0.310
(1, 4096, 14336)	0.098401	0.077046	0.134575	0.435717	1.277
(1, 14336, 4096)	0.084088	0.153705	0.134534	0.435533	0.547
(8, 4096, 4096)	0.028638	0.045108	0.050800	0.136783	0.635
(8, 4096, 1024)	0.015131	0.044912	0.025529	0.047025	0.337
(8, 4096, 14336)	0.099945	0.073439	0.135034	0.435975	1.364
(8, 14336, 4096)	0.083806	0.148976	0.134850	0.435791	0.563
(16, 4096, 4096)	0.029405	0.045662	0.050996	0.136979	0.642
(16, 4096, 1024)	0.020930	0.045057	0.025628	0.047123	0.463
(16, 4096, 14336)	0.100186	0.073811	0.135558	0.436499	1.365
(16, 14336, 4096)	0.084441	0.150410	0.135210	0.436151	0.561

[msaroufim]

Really not a fan of checking in notebooks, mind turning this into a python file, there should be a button to do this in jupyter, alternatively link to the tutorial hosted on google colab

[jcaip]

Sure, that fine with me.

[msaroufim]

we already have timing utilities in torchao/utils.py