chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,275 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Convert LLMCompressor compressed-tensors quantized model → MLX-native format.
|
||||
|
||||
Supports any bit-width (W4A16, W8A16, etc.) produced by LLMCompressor's
|
||||
pack-quantized format, including per-channel (strategy=channel) and per-group
|
||||
(strategy=group) quantization. Handles cross-shard layers automatically.
|
||||
|
||||
Usage:
|
||||
python convert_compressed_tensors_to_mlx.py /path/to/src /path/to/dst [--verify]
|
||||
"""
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import glob
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
|
||||
_MLX_MAX_GROUP_SIZE = {4: 128, 8: 128}
|
||||
_MLX_DEFAULT_GROUP_SIZE = 128
|
||||
|
||||
|
||||
def _effective_group_size(num_bits, config_group_size, in_features):
|
||||
if config_group_size is not None and config_group_size > 0:
|
||||
return config_group_size
|
||||
max_gs = _MLX_MAX_GROUP_SIZE.get(num_bits, _MLX_DEFAULT_GROUP_SIZE)
|
||||
return min(max_gs, in_features)
|
||||
|
||||
|
||||
def convert(src_dir: str, dst_dir: str) -> dict:
|
||||
from safetensors.torch import load_file, save_file
|
||||
import torch
|
||||
|
||||
src_dir = Path(src_dir)
|
||||
dst_dir = Path(dst_dir)
|
||||
dst_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with open(src_dir / "config.json") as f:
|
||||
config = json.load(f)
|
||||
|
||||
qconfig = config.get("quantization_config", {})
|
||||
group_cfg = qconfig.get("config_groups", {}).get("group_0", {}).get("weights", {})
|
||||
config_group_size = group_cfg.get("group_size", None)
|
||||
num_bits = group_cfg.get("num_bits", 4)
|
||||
symmetric = group_cfg.get("symmetric", True)
|
||||
strategy = group_cfg.get("strategy", "group")
|
||||
ignore = set(qconfig.get("ignore", []))
|
||||
|
||||
pack_factor = 32 // num_bits
|
||||
offset = 1 << (num_bits - 1)
|
||||
per_channel = (strategy == "channel") or (config_group_size is None)
|
||||
|
||||
print(f"Config: {num_bits}-bit, strategy={strategy}, "
|
||||
f"config_group_size={config_group_size}, "
|
||||
f"symmetric={symmetric}, offset={offset}, ignore={ignore}")
|
||||
|
||||
# Copy auxiliary files
|
||||
skip_ext = {".safetensors"}
|
||||
skip_names = {"model.safetensors.index.json", "config.json"}
|
||||
for f in src_dir.iterdir():
|
||||
if f.suffix in skip_ext or f.name in skip_names:
|
||||
continue
|
||||
if f.is_file():
|
||||
shutil.copy2(f, dst_dir / f.name)
|
||||
|
||||
# ── Load ALL shards, track which shard each key came from ────
|
||||
st_files = sorted(glob.glob(str(src_dir / "model*.safetensors")))
|
||||
if not st_files:
|
||||
st_files = sorted(glob.glob(str(src_dir / "*.safetensors")))
|
||||
|
||||
# key → (tensor, source_filename)
|
||||
all_tensors = {}
|
||||
for st_file in st_files:
|
||||
fname = os.path.basename(st_file)
|
||||
print(f" Loading {fname}...")
|
||||
for key, tensor in load_file(st_file).items():
|
||||
all_tensors[key] = (tensor, fname)
|
||||
|
||||
# ── Classify: quantized vs plain ─────────────────────────────
|
||||
quant_layers = {} # base → {packed, scale, shape}
|
||||
plain = {} # key → (tensor, fname)
|
||||
|
||||
for key, (tensor, fname) in all_tensors.items():
|
||||
matched = False
|
||||
for suffix, part in [(".weight_packed", "packed"),
|
||||
(".weight_scale", "scale"),
|
||||
(".weight_shape", "shape")]:
|
||||
if key.endswith(suffix):
|
||||
base = key.removesuffix(suffix)
|
||||
quant_layers.setdefault(base, {})[part] = (tensor, fname)
|
||||
matched = True
|
||||
break
|
||||
if not matched:
|
||||
plain[key] = (tensor, fname)
|
||||
|
||||
# ── Convert ──────────────────────────────────────────────────
|
||||
# Output: per-shard tensors
|
||||
import torch
|
||||
shard_out = {} # fname → {key: tensor}
|
||||
weight_map = {}
|
||||
total_converted = 0
|
||||
actual_group_size = None
|
||||
|
||||
for base, parts in sorted(quant_layers.items()):
|
||||
if "packed" not in parts:
|
||||
print(f" ⚠ Skipping {base}: no weight_packed")
|
||||
continue
|
||||
if "scale" not in parts:
|
||||
print(f" ⚠ Skipping {base}: no weight_scale")
|
||||
continue
|
||||
|
||||
packed, packed_fname = parts["packed"]
|
||||
scale, _ = parts["scale"]
|
||||
shape_tensor = parts.get("shape", (None, None))[0]
|
||||
|
||||
if shape_tensor is not None:
|
||||
in_features = int(shape_tensor[1].item())
|
||||
out_features = int(shape_tensor[0].item())
|
||||
else:
|
||||
in_features = packed.shape[1] * pack_factor
|
||||
out_features = packed.shape[0]
|
||||
|
||||
gs = _effective_group_size(num_bits, config_group_size, in_features)
|
||||
if actual_group_size is None:
|
||||
actual_group_size = gs
|
||||
print(f" Using group_size={gs} "
|
||||
f"({'broadcast from per-channel' if per_channel else 'from config'})")
|
||||
|
||||
# packed → uint32 as-is
|
||||
mlx_weight = torch.from_numpy(packed.numpy().view(np.uint32).copy())
|
||||
|
||||
# scale → fp16, broadcast if per-channel
|
||||
scale_fp16 = scale.to(torch.float16)
|
||||
n_groups = in_features // gs
|
||||
if scale_fp16.shape[1] < n_groups:
|
||||
scale_fp16 = scale_fp16.expand(out_features, n_groups).contiguous()
|
||||
|
||||
mlx_scales = scale_fp16
|
||||
mlx_biases = -float(offset) * mlx_scales
|
||||
|
||||
# Write to the shard that owns weight_packed
|
||||
shard_out.setdefault(packed_fname, {})[f"{base}.weight"] = mlx_weight
|
||||
shard_out[packed_fname][f"{base}.scales"] = mlx_scales
|
||||
shard_out[packed_fname][f"{base}.biases"] = mlx_biases
|
||||
total_converted += 1
|
||||
|
||||
# Plain tensors
|
||||
total_passthrough = 0
|
||||
for key, (tensor, fname) in plain.items():
|
||||
if tensor.dtype == torch.bfloat16:
|
||||
tensor = tensor.to(torch.float16)
|
||||
shard_out.setdefault(fname, {})[key] = tensor
|
||||
total_passthrough += 1
|
||||
|
||||
# ── Save shards ──────────────────────────────────────────────
|
||||
for fname, tensors in sorted(shard_out.items()):
|
||||
out_path = dst_dir / fname
|
||||
save_file(tensors, str(out_path))
|
||||
print(f" Saved {fname}: {len(tensors)} tensors")
|
||||
for key in tensors:
|
||||
weight_map[key] = fname
|
||||
|
||||
# Index
|
||||
with open(dst_dir / "model.safetensors.index.json", "w") as f:
|
||||
json.dump({"metadata": {"total_size": 0}, "weight_map": weight_map}, f, indent=2)
|
||||
|
||||
# Config
|
||||
new_config = dict(config)
|
||||
new_config.pop("quantization_config", None)
|
||||
new_config["quantization"] = {
|
||||
"group_size": actual_group_size or _MLX_DEFAULT_GROUP_SIZE,
|
||||
"bits": num_bits,
|
||||
}
|
||||
with open(dst_dir / "config.json", "w") as f:
|
||||
json.dump(new_config, f, indent=2)
|
||||
|
||||
stats = {"bits": num_bits, "group_size": actual_group_size,
|
||||
"strategy": strategy, "converted": total_converted,
|
||||
"passthrough": total_passthrough}
|
||||
print(f"\n✅ Done — {total_converted} quantized layers converted.")
|
||||
print(f" Output: {dst_dir}")
|
||||
return stats
|
||||
|
||||
|
||||
def verify(src_dir: str, dst_dir: str, n_check: int = 3):
|
||||
from safetensors.torch import load_file
|
||||
import torch
|
||||
|
||||
try:
|
||||
import mlx.core as mx
|
||||
except ImportError:
|
||||
print("⚠ mlx not available — skipping verification")
|
||||
return
|
||||
|
||||
src_dir, dst_dir = Path(src_dir), Path(dst_dir)
|
||||
|
||||
with open(src_dir / "config.json") as f:
|
||||
src_cfg = json.load(f)
|
||||
qconfig = src_cfg.get("quantization_config", {})
|
||||
group_cfg = qconfig.get("config_groups", {}).get("group_0", {}).get("weights", {})
|
||||
config_group_size = group_cfg.get("group_size", None)
|
||||
num_bits = group_cfg.get("num_bits", 4)
|
||||
|
||||
with open(dst_dir / "config.json") as f:
|
||||
dst_cfg = json.load(f).get("quantization", {})
|
||||
mlx_group_size = dst_cfg.get("group_size", 128)
|
||||
bits = dst_cfg.get("bits", num_bits)
|
||||
offset = 1 << (bits - 1)
|
||||
pack_factor = 32 // bits
|
||||
|
||||
# Load all shards from both dirs
|
||||
src_all = {}
|
||||
for f in sorted(glob.glob(str(src_dir / "model*.safetensors"))):
|
||||
src_all.update(load_file(f))
|
||||
dst_all = {}
|
||||
for f in sorted(glob.glob(str(dst_dir / "model*.safetensors"))):
|
||||
dst_all.update(load_file(f))
|
||||
|
||||
checked = 0
|
||||
for key in sorted(src_all):
|
||||
if not key.endswith(".weight_packed") or checked >= n_check:
|
||||
continue
|
||||
base = key.removesuffix(".weight_packed")
|
||||
|
||||
packed = src_all[f"{base}.weight_packed"].numpy().view(np.uint32)
|
||||
scale = src_all[f"{base}.weight_scale"].to(torch.float32).numpy()
|
||||
shape_t = src_all.get(f"{base}.weight_shape")
|
||||
if shape_t is not None:
|
||||
out_feat, in_feat = int(shape_t[0].item()), int(shape_t[1].item())
|
||||
else:
|
||||
in_feat = packed.shape[1] * pack_factor
|
||||
out_feat = packed.shape[0]
|
||||
|
||||
ref = np.zeros((out_feat, in_feat), dtype=np.float32)
|
||||
for i in range(pack_factor):
|
||||
vals = (packed >> (bits * i)) & ((1 << bits) - 1)
|
||||
signed = vals.astype(np.int32) - offset
|
||||
cols = np.arange(packed.shape[1]) * pack_factor + i
|
||||
cols = cols[cols < in_feat]
|
||||
if scale.shape[1] == 1:
|
||||
ref[:, cols] = signed[:, :len(cols)] * scale[:, 0:1]
|
||||
else:
|
||||
groups = cols // config_group_size
|
||||
ref[:, cols] = signed[:, :len(cols)] * scale[:, groups]
|
||||
|
||||
w = mx.array(dst_all[f"{base}.weight"].numpy().view(np.uint32))
|
||||
s = mx.array(dst_all[f"{base}.scales"].to(torch.float32).numpy())
|
||||
b = mx.array(dst_all[f"{base}.biases"].to(torch.float32).numpy())
|
||||
mlx_deq = np.array(mx.dequantize(w, s, b, group_size=mlx_group_size, bits=bits))
|
||||
|
||||
max_diff = np.max(np.abs(ref - mlx_deq))
|
||||
cos = (np.dot(ref.flatten(), mlx_deq.flatten())
|
||||
/ (np.linalg.norm(ref.flatten()) * np.linalg.norm(mlx_deq.flatten()) + 1e-12))
|
||||
status = "✅" if cos > 0.9999 else "❌"
|
||||
print(f" {status} {base}: max_diff={max_diff:.2e}, cos={cos:.8f}")
|
||||
checked += 1
|
||||
|
||||
print(f"\n Verified {checked} layers")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Convert LLMCompressor compressed-tensors → MLX format"
|
||||
)
|
||||
parser.add_argument("src", help="Source compressed-tensors model directory")
|
||||
parser.add_argument("dst", help="Destination MLX model directory")
|
||||
parser.add_argument("--verify", action="store_true",
|
||||
help="Spot-check dequantized values after conversion")
|
||||
args = parser.parse_args()
|
||||
|
||||
convert(args.src, args.dst)
|
||||
if args.verify:
|
||||
print("\n--- Verification ---")
|
||||
verify(args.src, args.dst)
|
||||
@@ -0,0 +1,143 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Unified PPL evaluation: FP16, W8A16, Cider per-channel, per-group gs=64, per-group gs=128."""
|
||||
import argparse, math, sys, time, gc, os
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
import numpy as np
|
||||
import tqdm
|
||||
from mlx_lm.utils import get_total_parameters, load
|
||||
from cider.nn import CiderLinear
|
||||
|
||||
home_dir = os.path.expanduser("~")
|
||||
WIKITEXT_LOCAL = os.path.join(home_dir, "Downloads/wikitext/wikitext-2-raw-v1")
|
||||
|
||||
FP16_MODEL = os.path.join(home_dir, "Downloads/Meta-Llama-3-8B")
|
||||
W8A16_MODEL = os.path.join(home_dir, "Downloads/Llama3-8B-GPTQ-W8A16-mlx")
|
||||
|
||||
def load_wikitext_local(tokenizer, split="test"):
|
||||
from datasets import load_dataset
|
||||
split_map = {
|
||||
"test": f"{WIKITEXT_LOCAL}/test-00000-of-00001.parquet",
|
||||
}
|
||||
data = load_dataset("parquet", data_files={split: split_map[split]}, split=split)
|
||||
encodings = tokenizer.encode("\n\n".join(data["text"]), return_tensors="pt")
|
||||
return encodings.numpy()
|
||||
|
||||
def convert_model_with_gs(model, target_group_size):
|
||||
counter = [0]
|
||||
_TARGET = (nn.Linear, nn.QuantizedLinear)
|
||||
def _walk(module):
|
||||
for name, child in module.children().items():
|
||||
if isinstance(child, _TARGET):
|
||||
setattr(module, name, CiderLinear.from_float(child, target_group_size=target_group_size))
|
||||
counter[0] += 1
|
||||
if counter[0] % 28 == 0: gc.collect()
|
||||
elif isinstance(child, list):
|
||||
for i, item in enumerate(child):
|
||||
if isinstance(item, _TARGET):
|
||||
child[i] = CiderLinear.from_float(item, target_group_size=target_group_size)
|
||||
counter[0] += 1
|
||||
if counter[0] % 28 == 0: gc.collect()
|
||||
elif isinstance(item, nn.Module):
|
||||
_walk(item)
|
||||
elif isinstance(child, nn.Module):
|
||||
_walk(child)
|
||||
_walk(model)
|
||||
return counter[0]
|
||||
|
||||
def eval_ppl(model, data, n_samples, seq_length):
|
||||
all_losses = []
|
||||
total_chunks = data.shape[1] // seq_length
|
||||
n = n_samples if n_samples > 0 else total_chunks
|
||||
for i in tqdm.tqdm(range(n), desc="Evaluating PPL"):
|
||||
batch = mx.array(data[:, i * seq_length: (i + 1) * seq_length])
|
||||
logits = model(batch[:, :-1]).astype(mx.float32)
|
||||
losses = nn.losses.cross_entropy(logits, batch[:, 1:], reduction="none")
|
||||
mx.eval(losses)
|
||||
all_losses.append(losses.flatten())
|
||||
all_losses = mx.concatenate(all_losses)
|
||||
mean_loss = all_losses.mean().item()
|
||||
ppl = math.exp(mean_loss)
|
||||
std_dev = mx.sqrt(mx.var(all_losses, ddof=1)).item()
|
||||
se_ppl = ppl * (std_dev / math.sqrt(all_losses.size))
|
||||
return ppl, se_ppl
|
||||
|
||||
def run_config(config_name, model_path, cider_gs=None, data=None, n_samples=50, seq_length=2048):
|
||||
"""Run a single config and return results dict."""
|
||||
print(f"\n{'='*60}")
|
||||
print(f" {config_name}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
print(f" Loading {model_path}...")
|
||||
model, tokenizer = load(model_path, tokenizer_config={"trust_remote_code": True})
|
||||
|
||||
if cider_gs is not None:
|
||||
gs_str = "per-channel" if cider_gs == 0 else f"per-group(gs={cider_gs})"
|
||||
print(f" Converting to CiderLinear ({gs_str})...")
|
||||
t0 = time.perf_counter()
|
||||
n = convert_model_with_gs(model, cider_gs)
|
||||
print(f" {n} layers converted in {time.perf_counter()-t0:.1f}s")
|
||||
|
||||
# Warmup
|
||||
dummy = mx.array([[1, 2, 3, 4, 5, 6, 7, 8]])
|
||||
_ = model(dummy); mx.eval(_)
|
||||
|
||||
if data is None:
|
||||
data = load_wikitext_local(tokenizer, "test")
|
||||
|
||||
mx.reset_peak_memory()
|
||||
start = time.time()
|
||||
ppl, se = eval_ppl(model, data, n_samples, seq_length)
|
||||
elapsed = time.time() - start
|
||||
peak_mem = mx.get_peak_memory() / 1e9
|
||||
|
||||
result = {"name": config_name, "ppl": ppl, "se": se, "time": elapsed, "mem": peak_mem}
|
||||
print(f" PPL = {ppl:.3f} ± {se:.3f} | {elapsed:.1f}s | {peak_mem:.2f} GB")
|
||||
|
||||
# Free memory
|
||||
del model
|
||||
gc.collect()
|
||||
mx.metal.clear_cache()
|
||||
|
||||
return result, data
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--num-samples", type=int, default=-1)
|
||||
parser.add_argument("--sequence-length", type=int, default=2048)
|
||||
args = parser.parse_args()
|
||||
|
||||
np.random.seed(123)
|
||||
mx.random.seed(123)
|
||||
|
||||
configs = [
|
||||
# ("Baseline FP16", FP16_MODEL, None),
|
||||
# ("Baseline W8A16 (MLX native)", W8A16_MODEL, None),
|
||||
# ("Cider W8A8 per-channel", W8A16_MODEL, 0),
|
||||
("Cider W8A8 per-group(gs=64)", W8A16_MODEL, 64),
|
||||
("Cider W8A8 per-group(gs=128)", W8A16_MODEL, 128),
|
||||
]
|
||||
|
||||
results = []
|
||||
data = None
|
||||
|
||||
for name, model_path, cider_gs in configs:
|
||||
np.random.seed(123)
|
||||
mx.random.seed(123)
|
||||
r, data = run_config(name, model_path, cider_gs, data, args.num_samples, args.sequence_length)
|
||||
results.append(r)
|
||||
|
||||
# Summary table
|
||||
print(f"\n\n{'='*70}")
|
||||
print(f" PPL COMPARISON SUMMARY (Qwen3-8B, wikitext-2, seq={args.sequence_length}, n={args.num_samples})")
|
||||
print(f"{'='*70}")
|
||||
print(f" {'Config':<30} {'PPL':>10} {'±SE':>8} {'Time':>8} {'Mem':>8}")
|
||||
print(f" {'-'*30} {'-'*10} {'-'*8} {'-'*8} {'-'*8}")
|
||||
|
||||
baseline_ppl = results[0]["ppl"] if results else None
|
||||
for r in results:
|
||||
delta = f"(+{(r['ppl']-baseline_ppl)/baseline_ppl*100:.2f}%)" if baseline_ppl and r['ppl'] != baseline_ppl else ""
|
||||
print(f" {r['name']:<30} {r['ppl']:>8.3f} ±{r['se']:.3f} {r['time']:>6.1f}s {r['mem']:>6.2f}GB {delta}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,118 @@
|
||||
import base64
|
||||
from io import BytesIO
|
||||
import json
|
||||
from datasets import Dataset
|
||||
import torch
|
||||
from compressed_tensors.offload import dispatch_model
|
||||
from datasets import load_dataset
|
||||
|
||||
# Note: this is an optional utility for processing vision inputs for qwen.
|
||||
# This can be installed via the "qwen" extra
|
||||
from qwen_vl_utils import process_vision_info
|
||||
from transformers import AutoProcessor, Qwen3VLForConditionalGeneration
|
||||
|
||||
from llmcompressor import oneshot
|
||||
from llmcompressor.modifiers.quantization import QuantizationModifier
|
||||
from llmcompressor.modifiers.awq import AWQModifier
|
||||
import random
|
||||
# Load model.
|
||||
model_id = "/mnt/data/ws/project/ms-swift/output/v3_sft_4b/v0-20260421-095015/checkpoint-50349"
|
||||
model = Qwen3VLForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
|
||||
processor = AutoProcessor.from_pretrained(model_id)
|
||||
|
||||
# Oneshot arguments
|
||||
DATASET_ID = "/mnt/data/ws/project/ms-swift/train_v4.jsonl"
|
||||
DATASET_SPLIT = "test[:512]"
|
||||
NUM_CALIBRATION_SAMPLES = 512
|
||||
MAX_SEQUENCE_LENGTH = 40000
|
||||
|
||||
# Load dataset and preprocess.
|
||||
def reservoir_sampling(file_path, k):
|
||||
"""
|
||||
从大文件中均匀随机采样 k 个样本
|
||||
|
||||
时间复杂度: O(n)
|
||||
空间复杂度: O(k)
|
||||
"""
|
||||
reservoir = []
|
||||
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
# 先填充前 k 个样本
|
||||
for i, line in enumerate(f):
|
||||
l = line.strip()
|
||||
if len(l) == 0:
|
||||
continue
|
||||
if i < k:
|
||||
reservoir.append(l)
|
||||
else:
|
||||
# 对于第 i 个样本，以 k/i 的概率替换蓄水池中的某个样本
|
||||
j = random.randint(0, i)
|
||||
if j < k:
|
||||
reservoir[j] = l
|
||||
|
||||
return reservoir
|
||||
|
||||
|
||||
# Apply chat template and tokenize inputs.
|
||||
calib_data_raw = reservoir_sampling(DATASET_ID, NUM_CALIBRATION_SAMPLES)
|
||||
|
||||
|
||||
def preprocess_and_tokenize(example):
|
||||
# preprocess
|
||||
|
||||
text = processor.apply_chat_template(
|
||||
example["messages"], tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
if "<image>" not in text:
|
||||
return None
|
||||
text = text.replace("<image>", "<|vision_start|><|image_pad|><|vision_end|>")
|
||||
|
||||
res = processor(
|
||||
text=[text],
|
||||
images=example["images"],
|
||||
padding=False,
|
||||
max_length=MAX_SEQUENCE_LENGTH,
|
||||
truncation=True,
|
||||
)
|
||||
# tokenize
|
||||
return res
|
||||
|
||||
calib_data = []
|
||||
for i, item in enumerate(calib_data_raw):
|
||||
print("load %d/%d items..." % (i, len(calib_data_raw)))
|
||||
u = preprocess_and_tokenize(json.loads(item))
|
||||
if u is None:
|
||||
continue
|
||||
calib_data.append(preprocess_and_tokenize(json.loads(item)))
|
||||
|
||||
calib_dataset = Dataset.from_list(calib_data)
|
||||
# Define a oneshot data collator for multimodal inputs.
|
||||
def data_collator(batch):
|
||||
assert len(batch) == 1
|
||||
return {key: torch.tensor(value) for key, value in batch[0].items()}
|
||||
|
||||
|
||||
# Recipe
|
||||
recipe = [
|
||||
AWQModifier(duo_scaling=False),
|
||||
QuantizationModifier(
|
||||
scheme="W8A16",
|
||||
ignore=["re:.*lm_head", "re:.*visual.*"],
|
||||
),
|
||||
]
|
||||
# Perform oneshot
|
||||
oneshot(
|
||||
model=model,
|
||||
tokenizer=model_id,
|
||||
dataset=calib_dataset,
|
||||
recipe=recipe,
|
||||
max_seq_length=MAX_SEQUENCE_LENGTH,
|
||||
num_calibration_samples=NUM_CALIBRATION_SAMPLES,
|
||||
data_collator=data_collator,
|
||||
sequential_targets=["Qwen3VLTextDecoderLayer"],
|
||||
)
|
||||
|
||||
# Save to disk compressed.
|
||||
SAVE_DIR = model_id.rstrip("/").split("/")[-1] + "-W8A16"
|
||||
model.save_pretrained(SAVE_DIR, save_compressed=True)
|
||||
processor.save_pretrained(SAVE_DIR)
|
||||
Reference in New Issue
Block a user