Loading issue on ROCm — asymmetric `compressed-tensors` MoE not supported in vLLM
#2
by tensornet - opened
The model works great on CUDA, but on AMD ROCm (gfx1201 / RDNA4) with vLLM 0.27.1, loading fails with:
AssertionError: Only symmetric quantization is supported for MoE
Root cause: This model uses the compressed-tensors quantization format with symmetric: false (asymmetric quantization with zero-points). vLLM has a hard assertion blocking asymmetric compressed-tensors MoE on the ROCm Triton kernel path.
For context, vLLM merged asymmetric compressed-tensors MoE support for the CUDA Marlin kernel in PR #44025 (v0.23.0+), but the ROCm Triton path hasn't been patched with the equivalent. Your Qwen3.6-35B-A3B-AWQ-4bit model works fine on ROCm because it uses native awq format, which vLLM's ROCm loader handles natively.
What would help:
- A symmetric variant — Releasing a version with
symmetric: truein the quantization config would load on ROCm with zero changes. - A native
awq-format version — Like your (Qwen3.6 model) — that quantization path is fully supported on ROCm.
Just flagging in case others run into the same wall, and in case a symmetric/native-AWQ variant is feasible. Happy to test anything!
Thanks for flagging this with me. Please download the model into a folder and use the following script to convert into AutoAWQ format:
"""
Convert a compressed-tensors `pack-quantized` checkpoint into AutoAWQ format.
Source layout (per quantized Linear, 4-bit example):
.weight_packed (out, in // 8) int32 sequential LSB pack along axis=1
.weight_scale (out, G) fp* G = in // group_size
.weight_zero_point (out // 8, G) int32 sequential LSB pack along axis=0 (asym only)
.weight_shape (2,) int stored original (out, in) shape
Target layout (per quantized Linear):
.qweight (in, out // 8) int32 AWQ_ORDER pack along output axis
.scales (G, out) fp16
.qzeros (G, out // 8) int32 AWQ_ORDER pack along output axis
(filled with 2^(b-1) when source was symmetric)
Both formats use signed-then-shifted unsigned packing; the dequant identity
(q_signed - zp_signed) * scale == (q_unsigned - zp_unsigned) * scale
makes the relayout numerically exact (no requantization).
Only 4-bit, group strategy, single config_groups, no actorder is supported.
"""
import json
import shutil
from pathlib import Path
import torch
from safetensors import safe_open
from safetensors.torch import save_file
try:
from tqdm import tqdm
except ImportError:
def tqdm(it, **kwargs): # minimal fallback
desc = kwargs.get("desc", "")
items = list(it)
if desc:
print(f"{desc}: {len(items)} items")
return items
AWQ_ORDER = [0, 2, 4, 6, 1, 3, 5, 7]
def unpack_lsb_unsigned(packed: torch.Tensor, bits: int, axis: int, original_len: int) -> torch.Tensor:
"""Inverse of compressed-tensors `pack_to_int32` but returns unsigned values
in [0, 2^bits - 1] (skipping the final -offset signed cast).
"""
pf = 32 // bits
if 32 % bits != 0:
raise ValueError(f"unpack_lsb_unsigned only supports bits dividing 32 (got {bits})")
mask = (1 << bits) - 1
p = packed.to(torch.int32).movedim(axis, -1).contiguous()
*rest, K_packed = p.shape
shifts = (torch.arange(pf, dtype=torch.int32, device=packed.device) * bits).view(
*([1] * len(rest)), 1, pf
)
expanded = (p.unsqueeze(-1) >> shifts) & mask # (..., K_packed, pf)
expanded = expanded.view(*rest, K_packed * pf) # interleave back
expanded = expanded[..., :original_len] # drop padding
return expanded.movedim(-1, axis).contiguous()
def pack_awq(t_uint: torch.Tensor, bits: int) -> torch.Tensor:
"""Pack unsigned int values along the LAST axis using AutoAWQ's interleave.
Mirrors `pack_exllama` from AutoAWQ but with the AWQ_ORDER lane permutation
that AutoAWQ kernels expect.
"""
pf = 32 // bits
if pf != 8:
raise NotImplementedError(f"pack_awq only implemented for 4-bit (pf=8), got bits={bits}")
*rest, M = t_uint.shape
if M % pf != 0:
raise ValueError(f"last dim {M} not divisible by pack_factor {pf}")
t = t_uint.to(torch.int32) & ((1 << bits) - 1)
t = t.view(*rest, M // pf, pf)
order = torch.tensor(AWQ_ORDER, dtype=torch.long, device=t.device)
t = t.index_select(dim=-1, index=order) # apply AWQ_ORDER
shifts = (torch.arange(pf, dtype=torch.int32, device=t.device) * bits)
shifts = shifts.view(*([1] * (t.ndim - 1)), pf)
packed = (t << shifts).sum(dim=-1, dtype=torch.int32) # (..., M // pf)
return packed.contiguous()
def convert_linear(
weight_packed: torch.Tensor, # (out, in // pf) int32
weight_scale: torch.Tensor, # (out, G)
weight_zero_point: torch.Tensor | None, # (out // pf, G) int32, asym only
bits: int,
group_size: int,
out_features: int,
in_features: int,
symmetric: bool,
device: torch.device,
) -> dict:
pf = 32 // bits
G = in_features // group_size
if in_features % group_size != 0:
raise ValueError(f"in_features={in_features} not divisible by group_size={group_size}")
if out_features % pf != 0:
raise ValueError(f"out_features={out_features} not divisible by pack_factor={pf}")
if weight_scale.shape != (out_features, G):
raise ValueError(f"weight_scale shape {tuple(weight_scale.shape)} != (out={out_features}, G={G})")
# Move to compute device
weight_packed = weight_packed.to(device)
weight_scale = weight_scale.to(device)
# --- 1. Unpack source weights to unsigned uint values (out, in) in [0, 2^b - 1] ---
iweight_uint = unpack_lsb_unsigned(weight_packed, bits, axis=1, original_len=in_features)
if iweight_uint.shape != (out_features, in_features):
raise RuntimeError(
f"unpacked weight shape {tuple(iweight_uint.shape)} != (out={out_features}, in={in_features})"
)
# --- 2. Unsigned zero points (out, G) in [0, 2^b - 1] ---
if symmetric:
if weight_zero_point is not None:
raise RuntimeError("symmetric=True but weight_zero_point is present")
izeros_uint = torch.full((out_features, G), 1 << (bits - 1), dtype=torch.int32, device=device)
else:
if weight_zero_point is None:
raise RuntimeError("symmetric=False but weight_zero_point missing")
izeros_uint = unpack_lsb_unsigned(
weight_zero_point.to(device), bits, axis=0, original_len=out_features
)
if izeros_uint.shape != (out_features, G):
raise RuntimeError(
f"unpacked zp shape {tuple(izeros_uint.shape)} != (out={out_features}, G={G})"
)
# --- 3. Reorient to AWQ axes and pack with AWQ_ORDER along output dim ---
iweight_T = iweight_uint.t().contiguous() # (in, out)
izeros_T = izeros_uint.t().contiguous() # (G, out)
qweight = pack_awq(iweight_T, bits).cpu() # (in, out // pf)
qzeros = pack_awq(izeros_T, bits).cpu() # (G, out // pf)
scales = weight_scale.t().to(torch.float16).contiguous().cpu() # (G, out)
return {"qweight": qweight, "qzeros": qzeros, "scales": scales}
def convert_checkpoint(src_dir: str, dst_dir: str, version: str = "gemm", device: str = "auto") -> None:
src, dst = Path(src_dir), Path(dst_dir)
dst.mkdir(parents=True, exist_ok=True)
if device == "auto":
compute_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
else:
compute_device = torch.device(device)
print(f"Compute device: {compute_device}")
with open(src / "config.json") as f:
config = json.load(f)
qcfg = config.get("quantization_config")
if qcfg is None:
raise ValueError("source config.json has no quantization_config")
fmt = qcfg.get("format")
if fmt not in ("pack-quantized", "pack_quantized"):
raise ValueError(f"expected format=pack-quantized, got {fmt!r}")
groups = qcfg.get("config_groups") or {}
if len(groups) != 1:
raise NotImplementedError(f"multi-group schemes not supported (got {list(groups)})")
wcfg = next(iter(groups.values()))["weights"]
bits = wcfg["num_bits"]
group_size = wcfg["group_size"]
symmetric = bool(wcfg["symmetric"])
if wcfg.get("type", "int") != "int":
raise NotImplementedError(f"weights.type={wcfg['type']!r}, only 'int' is supported")
if wcfg.get("strategy") != "group":
raise NotImplementedError(f"strategy={wcfg.get('strategy')!r}, only 'group' is supported")
if wcfg.get("actorder") not in (None, "null", False, "static"):
raise NotImplementedError(f"actorder={wcfg.get('actorder')!r} not implemented (AWQ has no g_idx)")
if bits != 4:
raise NotImplementedError(f"only 4-bit conversion is implemented (got {bits})")
ignore_list = list(qcfg.get("ignore") or [])
shard_files = sorted(src.glob("model-*.safetensors"))
if not shard_files:
shard_files = [src / "model.safetensors"]
multi_shard = len(shard_files) > 1
weight_map: dict[str, str] = {}
total_size = 0
total_converted = 0
total_passthrough = 0
for shard_idx, shard in enumerate(tqdm(shard_files, desc="Shards")):
# Load one shard at a time — keeps peak RAM bounded to ~1 shard
shard_tensors: dict[str, torch.Tensor] = {}
with safe_open(str(shard), framework="pt") as f:
for k in f.keys():
shard_tensors[k] = f.get_tensor(k)
linear_prefixes = sorted(
k[: -len(".weight_packed")] for k in shard_tensors if k.endswith(".weight_packed")
)
shard_out: dict[str, torch.Tensor] = {}
for prefix in tqdm(linear_prefixes, desc=f" {shard.name}", leave=False):
wpk = f"{prefix}.weight_packed"
wsk = f"{prefix}.weight_scale"
wzk = f"{prefix}.weight_zero_point"
wshk = f"{prefix}.weight_shape"
wgk = f"{prefix}.weight_g_idx"
weight_packed = shard_tensors.pop(wpk)
weight_scale = shard_tensors.pop(wsk)
weight_zp = shard_tensors.pop(wzk, None)
weight_shape = shard_tensors.pop(wshk, None)
if wgk in shard_tensors:
raise NotImplementedError(f"{wgk} present — actorder/g_idx not supported")
if weight_shape is not None:
out_features, in_features = int(weight_shape[0].item()), int(weight_shape[1].item())
else:
out_features = weight_packed.shape[0]
in_features = weight_packed.shape[1] * (32 // bits)
packed = convert_linear(
weight_packed=weight_packed,
weight_scale=weight_scale,
weight_zero_point=weight_zp,
bits=bits,
group_size=group_size,
out_features=out_features,
in_features=in_features,
symmetric=symmetric,
device=compute_device,
)
for suffix, t in packed.items():
shard_out[f"{prefix}.{suffix}"] = t
total_converted += 1
# Carry over all remaining tensors in this shard unchanged
for k, t in shard_tensors.items():
shard_out[k] = t
total_passthrough += len(shard_tensors)
# Write this shard immediately, then free memory
if multi_shard:
out_name = f"model-{shard_idx + 1:05d}-of-{len(shard_files):05d}.safetensors"
else:
out_name = "model.safetensors"
out_path = dst / out_name
print(f" Writing {out_path.name} ({len(shard_out)} tensors)...")
save_file(shard_out, str(out_path))
for k, t in shard_out.items():
weight_map[k] = out_name
total_size += t.nbytes
del shard_tensors, shard_out
if compute_device.type == "cuda":
torch.cuda.empty_cache()
# Write index file for multi-shard output
if multi_shard:
index = {"metadata": {"total_size": total_size}, "weight_map": weight_map}
with open(dst / "model.safetensors.index.json", "w") as f:
json.dump(index, f, indent=2)
print(f"Converted {total_converted} linears; carried over {total_passthrough} tensors unchanged.")
# --- Rewrite config.json with AWQ schema ---
awq_quant_config = {
"quant_method": "awq",
"zero_point": not symmetric,
"group_size": group_size,
"bits": bits,
"version": version,
"modules_to_not_convert": ignore_list or None,
}
new_config = {k: v for k, v in config.items() if k != "quantization_config"}
new_config["quantization_config"] = awq_quant_config
with open(dst / "config.json", "w") as f:
json.dump(new_config, f, indent=2)
# --- Sidecar quant_config.json (legacy AutoAWQ readers) ---
sidecar = {
"zero_point": not symmetric,
"q_group_size": group_size,
"w_bit": bits,
"version": version,
"modules_to_not_convert": ignore_list or None,
}
with open(dst / "quant_config.json", "w") as f:
json.dump(sidecar, f, indent=2)
# --- Copy tokenizer/aux files ---
for name in (
"tokenizer.json",
"tokenizer.model",
"tokenizer_config.json",
"special_tokens_map.json",
"generation_config.json",
"chat_template.jinja",
"added_tokens.json",
"merges.txt",
"vocab.json",
):
p = src / name
if p.exists():
shutil.copy2(p, dst / name)
print(f"Wrote AWQ checkpoint to {dst}")
if __name__ == "__main__":
import argparse
ap = argparse.ArgumentParser()
ap.add_argument("--src", required=True, help="Source compressed-tensors pack-quantized model dir")
ap.add_argument("--dst", required=True, help="Destination AutoAWQ model dir")
ap.add_argument("--version", default="gemm", choices=["gemm", "gemv", "marlin"])
ap.add_argument(
"--device",
default="auto",
help="Compute device for tensor ops: 'auto' (default), 'cuda', 'cpu', 'cuda:1', etc.",
)
args = ap.parse_args()
convert_checkpoint(args.src, args.dst, version=args.version, device=args.device)