UniqAI/tree-of-thought-llm
1
0
Fork 0
mirror of https://github.com/princeton-nlp/tree-of-thought-llm synced 2025-05-30 01:48:21 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

adding ptq model support

Browse Source
This commit is contained in:
emilyworks 2024-12-04 00:26:31 -05:00
parent a0015eb0f7
commit 975a6e8608
1 changed files with 15 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
run.py
View File

@ -9,6 +9,7 @@ from src.tot.models import gpt_usage
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import torch.quantization
def run(args):
@ -17,14 +18,20 @@ def run(args):
'''
#load in non-gpt model in this driver function for now to avoid repeated loading later on
if args.backend == 'llama':
if not args.quantize:
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-3B-Instruct")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-3B-Instruct")
elif args.quantize == 'qat':
pass
# tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-3B-Instruct")
# model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-3B-Instruct-QLORA_INT4_EO8")
elif args.backend == 'spinquant':
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-3B-Instruct")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-3B-Instruct")
if args.quantize and args.quantize=='ptq':
model.train()
model.qconfig = torch.quantization.get_default_qconfig('x86')
torch.quantization.prepare(model, inplace=True)
for _, mod in model.named_modules():
if isinstance(mod, torch.nn.Embedding):
mod.qconfig = torch.ao.quantization.float_qparams_weight_only_qconfig
model = torch.quantization.convert(model, inplace=True)
model.load_state_dict(torch.load('quant_experiments/quantized_model.pth'))
model.eval()
elif args.backend == 'qat':
pass
# tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct")
# model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8")