| | --- |
| | library_name: transformers |
| | tags: [] |
| | --- |
| | |
| | ``` |
| | import torch |
| | from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig |
| | |
| | model_id = "facebook/opt-125m" |
| | |
| | from torchao.quantization import Float8DynamicActivationFloat8WeightConfig, PerRow |
| | quant_config = Float8DynamicActivationFloat8WeightConfig(granularity=PerRow()) |
| | |
| | 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}-float8dq-row-0.13-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") |
| | generated_ids = quantized_model.generate(**inputs, max_new_tokens=128) |
| | output_text = tokenizer.batch_decode( |
| | generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False |
| | ) |
| | print("Response:", output_text[0][len(prompt) :]) |
| | |
| | ``` |
