torchao-testing / opt-125m-Int4WeightOnlyConfig-v1-0.14.dev

model: opt-125m
config: Int4WeightOnlyConfig
config version: 1
torchao version: 0.14.dev

import torch
import io
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
from huggingface_hub import HfApi

model_id = "facebook/opt-125m"

from torchao.quantization import Int4WeightOnlyConfig
quant_config = Int4WeightOnlyConfig(group_size=128, version=1)

quantization_config = TorchAoConfig(quant_type=quant_config)
quantized_model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="cuda",
    torch_dtype=torch.bfloat16,
    quantization_config=quantization_config,
)
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Push to hub
USER_ID = "torchao-testing"
MODEL_NAME = model_id.split("/")[-1]
save_to = f"{USER_ID}/{MODEL_NAME}-Int4WeightOnlyConfig-v1-0.14.dev"
quantized_model.push_to_hub(save_to, safe_serialization=False)
tokenizer.push_to_hub(save_to)

# Manual Testing
prompt = "Hey, are you conscious? Can you talk to me?"
print("Prompt:", prompt)
inputs = tokenizer(
    prompt,
    return_tensors="pt",
).to("cuda")
# setting temperature to 0 to make sure result deterministic
generated_ids = quantized_model.generate(**inputs, max_new_tokens=128, temperature=0)
api = HfApi()
buf = io.BytesIO()
torch.save(prompt, buf)
api.upload_file(
    path_or_fileobj=buf,
    path_in_repo="model_prompt.pt",
    repo_id=save_to,
)

buf = io.BytesIO()
torch.save(generated_ids, buf)
api.upload_file(
    path_or_fileobj=buf,
    path_in_repo="model_output.pt",
    repo_id=save_to,
)

output_text = tokenizer.batch_decode(
    generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
)
print("Response:", output_text[0][len(prompt) :])