Update torchao.prototype.parq and add 4-bit Llama 3.2 1B benchmark

[lisjin]

We would like to merge recent changes from our open sourced library at https://github.com/facebookresearch/parq.

We have also benchmarked 4-bit Llama 3.2 1B fine-tuned for 25K steps on fineweb-edu using torchtune. We used PARQ's MaxUnifQuantizer and ProxHardQuant proximal mapping, which is equivalent to STE. Below are the relevant training config changes to the llama3_2/1B_full.yaml recipe.

batch_size: 8
epochs: 1
optimizer:
  _component_: torch.optim.AdamW
  lr: 4e-5
  weight_decay: 0.0
  betas: [0.9, 0.95]
  fused: True

lr_scheduler:
  _component_: torchtune.training.lr_schedulers.get_cosine_schedule_with_warmup
  num_warmup_steps: 2500

As shown in the table below, the resulting 4-bit model achieves well under 10% accuracy on most commonsense reasoning benchmarks relative to the pre-trained model.

Tasks	16-bit	4-bit	% diff
arc_challenge	0.3805	0.3575	-6.4
arc_easy	0.6309	0.6077	-3.7
hellaswag	0.6081	0.5423	-10.8
piqa	0.7410	0.7122	-3.9
winogrande	0.6022	0.5549	-7.8

[pytorch-bot]

🔗 Helpful Links

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

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

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

[andrewor14]

Hi @lisjin, thanks for the update! The results look great. As discussed offline, we usually only add submodules if other parts of torchao are using the submodules, which is not the case here. Do you mind making the latest changes to PARQ in the prototype version here in torchao instead?

[lisjin]

Do you mind making the latest changes to PARQ in the prototype version here in torchao instead?

Of course! I just removed the submodule and updated the prototype version instead.