[WIP] Apply SuperBlock to Llama

[mostafaelhoushi]

Still work in progress.
To run:

cd torchao/_models/llama

python generate.py --checkpoint_path ${CHECKPOINT_PATH}/model.pth --superblock

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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.

[jcaip]

@mostafaelhoushi Should this be changed to SupermaskReplacementClass?

[mostafaelhoushi]

hmmm... SupermaskReplacementClass constructor requires a lot of arguments like linear_sparsity, and linear_sp_tilesize, etc. How will we pass them here?

[mostafaelhoushi]

I think I need to do some more refactoring:

• the ViT benchmark code assumes that there is a model checkpoint trained with SuperBlock, and hence has SupermaskLinear layers and parameters
• the GPT-Fast code I wrote did a hack in which it converted Linear layers to SupermaskLinear layers then applied BSR sparsification.

[jcaip]

Ah, I think I read your changes wrong, I was assuming you created a SupermaskReplacementClass to combine SupermaskLinear, SupermaskConv, but I still you're still using those under the hood. I think this should be fine actually.

[jerryzh168]

should be torchao/prototype/sparsity ? #1013

[mostafaelhoushi]

I think this is because this PR forked from an old commit that had superblock in torchao/sparsity/prototype.
When finalizing the PR we can rebase on top of main and change the path of the directory.