Enhance Auto-Round

[yiliu30]

This PR includes several enhancements for Auto-round:

1. Bring torch.compile to speed up the Auto-round optimization process:

python torchao/prototype/autoround/autoround_llm.py -c

• meta-llama/Llama-2-7b-chat-hf, about 1.29x , 27 sec/block -> 21 sec/block
• meta-llama/Meta-Llama-3.1-8B-Instruct, about 1.23x, 32 sec/block -> 26 sec/block

yiliu30#18

2. Add AO_USE_DETERMINISTIC_ALGORITHMS for reproducing the lm-eval results:

AO_USE_DETERMINISTIC_ALGORITHMS=1 python torchao/prototype/autoround/eval_autoround.py

yiliu30#19

3. Expose gradient_accumulate_steps to users and update results:

• For meta-llama/Meta-Llama-3.1-8B-Instruct, here are the updated results:

	Avg.	Mmlu	Piqa	Winogrande	Hellaswag	Lambada_openai
bf16	0.7080	0.6783	0.8003	0.7403	0.5910	0.7303
torchao-int4wo	0.6883	0.6363	0.7938	0.7348	0.5784	0.6980
autoround-4bit	0.6996	0.6669	0.7916	0.7285	0.5846	0.7262
autoround-4bit*	0.7010	0.6621	0.7976	0.7316	0.5847	0.7291

For more details, please refer README.md

yiliu30#20

TODO:

• Add gradient_accumulate_steps and compile_optimization_process to eval.py after Move autoround from generate.py to eval.py #868 lands. Done by a5722d8

cc @wenhuach21 @thuang6

[pytorch-bot]

🔗 Helpful Links

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

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

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

[wenhuach21]

The reason why accumualte_gradient is better is:
to better align with the ao's API, we switched from random sampling in our original training implementation to using fixed samples within each batch, which reduces flexibility. However, gradient accumulation helps to partially recover this flexibility.

[yiliu30]

Hi @jerryzh168, could you please take a look, thanks!

[jerryzh168]

LGTM!