nm-testing/Qwen3-4B-mixed-quant-RTN-wNaM-a4

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

Layers cycle through symmetric and asymmetric variants of weight-only and weight+activation formats, with a symmetric W4A4 layer interleaved. W4A4asym is excluded because humming's dequant identity path (SourceType == TargetType) does not support zero_point.

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

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-a4" )

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) ("W7A8", 7, 8, True), ("W7A8asym", 7, 8, False), ("W7A8", 7, 8, True), ("W7A8asym", 7, 8, False), ("W7A8", 7, 8, True), ("W7A8asym", 7, 8, False), ("W7A8", 7, 8, True), ("W7A8asym", 7, 8, False),

("W3A4",       3, 4,    True),
("W3A4asym",   3, 4,    False),
("W4A4",       4, 4,    True),
("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),
("W8A8",       8, 8,    True),

("W7A8",       7, 8,    True),
("W7A8asym",   7, 8,    False),
("W7A8",       7, 8,    True),
("W7A8asym",   7, 8,    False),
("W7A8",       7, 8,    True),
("W7A8asym",   7, 8,    False),
("W7A8",       7, 8,    True),
("W7A8asym",   7, 8,    False),
("W7A8",       7, 8,    True),
("W7A8asym",   7, 8,    False),

]

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

config_groups = {} for i in range(num_layers): label, wbits, abits, sym = _BASE_FORMATS[ i % len(_BASE_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}")