Qwen3-4B-mixed-quant-RTN-wNaM-asym-w4a4

"""
Mixed-quantization checkpoint: W4A4 layer + symmetric AND asymmetric W5-W8.

Layer 0 gets W4A4 (symmetric); the remaining layers cycle through both
symmetric and asymmetric variants of weight-only and weight+activation
formats.

Formats cycled (by layer index):
    W5A16,      W5A16asym,
    W6A16,      W6A16asym,
    W7A16,      W7A16asym,
    W8A16,      W8A16asym,
    W5A8,       W5A8asym,
    W6A8,       W6A8asym,
    W7A8,       W7A8asym,

Usage:
    python mixed_quant_w4a4_asym.py
    python mixed_quant_w4a4_asym.py --model_id Qwen/Qwen3-4B
"""

import argparse
import os

from compressed_tensors.offload import dispatch_model
from compressed_tensors.quantization import (
    QuantizationArgs,
    QuantizationScheme,
    QuantizationStrategy,
    QuantizationType,
)
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier

parser = argparse.ArgumentParser()
parser.add_argument(
    "--model_id",
    type=str,
    default="Qwen/Qwen3-4B",
)
args = parser.parse_args()

SAVE_DIR = (
    args.model_id.rstrip("/").split("/")[-1] + "-mixed-quant-RTN-wnam-asym"
)

if os.path.exists(SAVE_DIR):
    print(f"Output already exists at {SAVE_DIR!r}, skipping.")
    exit(0)

_BASE_FORMATS = [
    # (label, weight_bits, act_bits_or_None, symmetric)
    ("W5A16",      5, None, True),
    ("W5A16asym",  5, None, False),
    ("W6A16",      6, None, True),
    ("W6A16asym",  6, None, False),
    ("W7A16",      7, None, True),
    ("W7A16asym",  7, None, False),
    ("W8A16",      8, None, True),
    ("W8A16asym",  8, None, False),
    ("W5A8",       5, 8,    True),
    ("W5A8asym",   5, 8,    False),
    ("W6A8",       6, 8,    True),
    ("W6A8asym",   6, 8,    False),
    ("W7A8",       7, 8,    True),
    ("W7A8asym",   7, 8,    False),
]

REMAINING_FORMATS = (
    _BASE_FORMATS
    + [("W4A4", 4, 4, True)]
    + _BASE_FORMATS
    + [("W4A4asym", 4, 4, False)]
)

num_layers = AutoConfig.from_pretrained(args.model_id).num_hidden_layers

config_groups = {}
for i in range(num_layers):
    label, wbits, abits, sym = REMAINING_FORMATS[
        i % len(REMAINING_FORMATS)
    ]

    weights = QuantizationArgs(
        num_bits=wbits,
        type=QuantizationType.INT,
        strategy=QuantizationStrategy.GROUP,
        group_size=128,
        symmetric=sym,
    )

    input_activations = None
    if abits is not None:
        input_activations = QuantizationArgs(
            num_bits=abits,
            type=QuantizationType.INT,
            strategy=QuantizationStrategy.TOKEN,
            dynamic=True,
            symmetric=True,
        )

    config_groups[f"layer_{i}_{label}"] = QuantizationScheme(
        targets=[f"re:model\\.layers\\.{i}\\..*_proj$"],
        weights=weights,
        input_activations=input_activations,
    )

    print(f"  layer {i:2d} -> {label}")

recipe = QuantizationModifier(
    config_groups=config_groups,
    ignore=["lm_head"],
)

model = AutoModelForCausalLM.from_pretrained(args.model_id, dtype="auto")
tokenizer = AutoTokenizer.from_pretrained(args.model_id)

oneshot(model=model, recipe=recipe)

print("\n\n========== SAMPLE GENERATION ==============")
dispatch_model(model)
input_ids = tokenizer(
    "Hello my name is", return_tensors="pt"
).input_ids.to(model.device)
output = model.generate(input_ids, max_new_tokens=100)
print(tokenizer.decode(output[0]))
print("==========================================\n\n")

model.save_pretrained(
    SAVE_DIR,
    save_compressed=True,
    quantization_format="pack-quantized",
)
tokenizer.save_pretrained(SAVE_DIR)
print(f"Saved to {SAVE_DIR}")