Views
No views yet
1from datasets import load_dataset
2from transformers import AutoTokenizer
3
4from auto_fp8 import AutoFP8ForCausalLM, BaseQuantizeConfig
5
6pretrained_model_dir = "Sao10K/MN-12B-Lyra-v3"
7quantized_model_dir = "MN-12B-Lyra-v3-FP8"
8
9tokenizer = AutoTokenizer.from_pretrained(pretrained_model_dir, use_fast=True, model_max_length=4096)
10tokenizer.pad_token = tokenizer.eos_token
11
12ds = load_dataset("mgoin/ultrachat_2k", split="train_sft").select(range(512))
13examples = [tokenizer.apply_chat_template(batch["messages"], tokenize=False) for batch in ds]
14examples = tokenizer(examples, padding=True, truncation=True, return_tensors="pt").to("cuda")
15
16quantize_config = BaseQuantizeConfig(
17 quant_method="fp8",
18 activation_scheme="static",
19 ignore_patterns=["re:.*lm_head"],
20)
21
22model = AutoFP8ForCausalLM.from_pretrained(
23 pretrained_model_dir, quantize_config=quantize_config
24)
25
26model.quantize(examples)
27model.save_quantized(quantized_model_dir)