Qwen3-235B-A22B-Instruct-2507-int4-mixed-AutoRound

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
from auto_round import AutoRound

model_name = "Qwen/Qwen3-235B-A22B-Instruct-2507"

model = AutoModelForCausalLM.from_pretrained(model_name,
                                             device_map="cpu", torch_dtype="auto")

tokenizer = AutoTokenizer.from_pretrained(model_name)

layer_config = {}
for n, m in model.named_modules():
    if "mlp.gate" in n: ## vllm only support 16 bit for this layer
        layer_config[n] = {"bits": 16}
    elif isinstance(m, torch.nn.Linear) and (not "expert" in n or "shared_experts" in n) and n != "lm_head":
        layer_config[n] = {"bits": 8, "group_size": 128}

autoround = AutoRound(model, tokenizer, iters=0, group_size=64, layer_config=layer_config)
output_dir = "/dataset/Qwen3-235B-A22B-Instruct-2507-int4-mixed"
autoround.quantize_and_save(output_dir)

## tricky code to handle qkv fusing issue, we will fix it in vllm later
import os
import json

config_path = os.path.join(output_dir, "config.json")

with open(config_path, "r") as file:
    config = json.load(file)
extra_config = config["quantization_config"]["extra_config"]
num_hidden_layers = config["num_hidden_layers"]
for i in range(num_hidden_layers):
    qkv_name = f"model.layers.{str(i)}.self_attn.qkv_proj"
    extra_config[qkv_name] = {"bits": 8, "group_size": 128}
with open(config_path, "w") as file:
    json.dump(config, file, indent=2)