Minachist's picture
Update quantize.py
9974f95 verified
Raw
History Blame Contribute Delete
2.24 kB
import torch
from transformers import AutoTokenizer
from auto_round import AutoRound
# You need to use datasets 3.6.0 since >=4.0 doesn't support codeparrot/github-code-clean dataset
# You also need torchvision for some reason even if it's not used, otherwise it will throw an error "Unrecognized image processor ...".
# TORCHVISION_AVAILABLE = False is there because datasets 3.6.0 uses the old torchvision api, and isn't compatible w/ this torchvision.'(And it will throw ImportError: cannot import name 'VideoReader' from 'torchvision.io')
# Downgrading torchvision might solve the issue, but there's also a chance that doing that breaks other deps. Since I don't want to damage my brain anymore, I decided to just simply told datasets not to use torchvision.
# To replicate the setup:
# uv venv --python 3.13 .venv
# source .venv/bin/activate
# uv pip install "torch==2.12.1" "torchvision==0.27.1" \
# --index-url https://download.pytorch.org/whl/cu130
# uv pip install "transformers==5.12.1" "auto-round-nightly==0.14.0.dev20260625" \
# "accelerate==1.14.0" "datasets==3.6.0" "pillow"
# This is why I don't like Python, it's dependency hell.
import datasets
datasets.config.TORCHVISION_AVAILABLE = False
model_name_or_path = "."
output_dir = "./Ornith-1.0-35B-INT8-AutoRound/"
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)
# auto-round doesn't support quantizing linear_attn, but I wanted to be sure.
ignore_keywords = [
"embed_tokens",
"linear_attn",
"shared_expert",
"mlp.gate",
"visual",
"lm_head",
]
layer_config = {}
for keyword in ignore_keywords:
layer_config[keyword] = {"bits": 16}
seqlen = 2048
nsamples = 1024
dataset = "NeelNanda/pile-10k:num=256,codeparrot/github-code-clean:num=768"
ar = AutoRound(
model=model_name_or_path,
tokenizer=tokenizer,
scheme="W8A16",
enable_torch_compile=False,
group_size=-1,
sym=True,
layer_config=layer_config,
dataset=dataset,
device_map="0,1",
batch_size=8,
seqlen=seqlen,
iters=1000,
nsamples=nsamples,
low_gpu_mem_usage=True,
)
ar.quantize_and_save(output_dir, format="auto_round")