AWQ does not actually Quant, outputs slightly bigger size file. - CoHere Model

[phaelon74]

Describe the bug
Using llm_compressor, successfully calling all libraries, and quantization finishing successfully, and yet the end file/s are even bigger than the input. Either:

1. It is failing to Quant the layers OR
2. It's failing to discard the original models weights when saving.

Expected behavior
The output directory is a properly 4bit quanted AWQ

Environment
Include all relevant environment information:

1. OS [e.g. Ubuntu 20.04]: Ubuntu 20.04
2. Python version [e.g. 3.7]:3.10.12
3. LLM Compressor version or commit hash [e.g. 0.1.0, f7245c8]: 0.5.1
4. ML framework version(s) [e.g. torch 2.3.1]: 2.5.1
5. Other Python package versions [e.g. vLLM, compressed-tensors, numpy, ONNX]:
6. Other relevant environment information [e.g. hardware, CUDA version]: CUDA 12.4

To Reproduce
I run my script that calls llm_compressor and does the quantization

Errors
No errors, the quantization finishes successfully

Additional context
Script here:

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer

from llmcompressor import oneshot
from llmcompressor.modifiers.awq import AWQModifier

# Select model and load it.
MODEL_ID = "./models/CohereLabs/c4ai-command-r7b-12-2024/"

model = AutoModelForCausalLM.from_pretrained(MODEL_ID,  device_map="auto", torch_dtype="auto")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, device_map="auto", trust_remote_code=True)

# Select calibration dataset.
DATASET_ID = "mit-han-lab/pile-val-backup"
DATASET_SPLIT = "validation"

# Select number of samples. 256 samples is a good place to start.
# Increasing the number of samples can improve accuracy.
NUM_CALIBRATION_SAMPLES = 256
MAX_SEQUENCE_LENGTH = 512

# Load dataset and preprocess.
ds = load_dataset(DATASET_ID, split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES}]")
ds = ds.shuffle(seed=42)


def preprocess(example):
    return {
        "text": tokenizer.apply_chat_template(
            [{"role": "user", "content": example["text"]}],
            tokenize=False,
        )
    }


ds = ds.map(preprocess)


# Tokenize inputs.
def tokenize(sample):
    return tokenizer(
        sample["text"],
        padding=False,
        max_length=MAX_SEQUENCE_LENGTH,
        truncation=True,
        add_special_tokens=False,
    )


# Configure the quantization algorithm to run.
recipe = [
    AWQModifier(
        model_type="custom",
        ignore=["lm_head"],
        mappings=[],
        module_mappers=[
            {
                "pattern": "re:.*.model.layers.\\d+$",
                "smooth_layer": "input_layernorm",
                "balance_layers": [
                    ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"),
                    ("mlp.down_proj", "mlp.gate_proj", "mlp.up_proj"),
                ],
                "to_quant_tensors": [
                    "self_attn.q_proj.weight",
                    "self_attn.k_proj.weight",
                    "self_attn.v_proj.weight",
                    "self_attn.o_proj.weight",
                    "mlp.gate_proj.weight",
                    "mlp.up_proj.weight",
                    "mlp.down_proj.weight",
                ],
            }
        ],
    )
]

# Apply algorithms.
oneshot(
    model=model,
    dataset=ds,
    recipe=recipe,
    max_seq_length=MAX_SEQUENCE_LENGTH,
    num_calibration_samples=NUM_CALIBRATION_SAMPLES,
)

# Confirm generations of the quantized model look sane.
print("\n\n")
print("========== SAMPLE GENERATION ==============")
input_ids = tokenizer("Hello my name is", return_tensors="pt").input_ids.to("cuda")
output = model.generate(input_ids, max_new_tokens=100)
print(tokenizer.decode(output[0]))
print("==========================================\n\n")

# Save to disk compressed.
SAVE_DIR = "./quantized/c4ai-command-r7b-4bpw-AWQ"
model.save_pretrained(SAVE_DIR, save_compressed=True)
tokenizer.save_pretrained(SAVE_DIR)

[brian-dellabetta]

Hi @phaelon74 , thanks for your interest in our AWQ feature! We have made a ton of improvements to it since the v0.5.1 release, including an update that will allow your recipe to work without having to explicitly add the QuantizationModifier. We plan to do a v0.6.0 release in the next week, in the meantime you can install from main branch or if you can't do that in your environments, see v0.5.1 example recipe here.

As for the mappings, with mappings=[] I don't expect it will do anything. Where did you see module_mappers? I would expect the initialization to fail with an unexpected constructor arg.

It looks like AutoAWQ has a different mapping configuration for these cohere models, than what we have for the model families we've mostly targeted (llama, qwen, mistral, phi, etc.). I'm not sure the cause of this discrepancy, will get in touch with our research team who knows the original maintainer. We have been using AutoAWQ as the main reference for our port.

[phaelon74]

Hi @phaelon74 , thanks for your interest in our AWQ feature! We have made a ton of improvements to it since the v0.5.1 release, including an update that will allow your recipe to work without having to explicitly add the QuantizationModifier. We plan to do a v0.6.0 release in the next week, in the meantime you can install from main branch or if you can't do that in your environments, see v0.5.1 example recipe here.

As for the mappings, with mappings=[] I don't expect it will do anything. Where did you see module_mappers? I would expect the initialization to fail with an unexpected constructor arg.

It looks like AutoAWQ has a different mapping configuration for these cohere models, than what we have for the model families we've mostly targeted (llama, qwen, mistral, phi, etc.). I'm not sure the cause of this discrepancy, will get in touch with our research team who knows the original maintainer. We have been using AutoAWQ as the main reference for our port.

It finished successfully, with module mappers, etc. It just saves the original model weights, so it's almost like it's failing to quantize and instead just saving the base model back lol.

I've done the oneshot, and had a VERY large script, but I'll take your example, and use it and not my edited one.

So 0.6.0 will have a lot of changes, making things much better. Does this include loading the models layers 1 at a time as opposed to all layers into VRAM?

[brian-dellabetta]

Great, lmk how it goes. I can look into adding the mappings specifically for the cohere models as well.

So 0.6.0 will have a lot of changes, making things much better. Does this include loading the models layers 1 at a time as opposed to all layers into VRAM?

Yes! That very recently landed on main, we wanted to get that in before the release:

• [Performance] Sequential onloading #1263

[phaelon74]

Great, lmk how it goes. I can look into adding the mappings specifically for the cohere models as well.

So 0.6.0 will have a lot of changes, making things much better. Does this include loading the models layers 1 at a time as opposed to all layers into VRAM?

Yes! That very recently landed on main, we wanted to get that in before the release:

* [[Performance] Sequential onloading #1263](https://github.com/vllm-project/llm-compressor/pull/1263)

Excellent. The script is running, I had to install lm_eval but now it's cooking. Once it finishes with llama-8b-instruct, proving it works, I'll pivot to Cohere again and try that model and see if I need to manually add mappings back in, as it relates to your script. My script most likely took parts and pieces from a lot of other scripts and I wonder if I had weird interlaced parts affecting it.

[phaelon74]

@brian-dellabetta I ran the example you provided verbatim, to make sure I wasn't skewing anything, on a runpod with four A40s and 200GB of memory. I am attempting to do an 8B AWQ and once calibration finishes, it consumes ALL system memory and I am no longer able to engage with the runpod. I have to kill it and start over:

Any ideas why the machine is consuming all the memory and then locking up?

[brian-dellabetta]

Hi @phaelon74 , just to be clear you are on llm-compressor 0.5.1 and trying to run the example file for v0.5.1 here -- https://github.com/vllm-project/llm-compressor/blob/0.5.1/examples/awq/awq_one_shot.py ? I was able to run this on a single 80GB A100, but AWQ in 0.5.1 is a super old version where we cache activations all at once, so memory requirements are pretty extreme. are you able to install from main? If not, we plan to kick off the 0.6.0 release soon

[brian-dellabetta]

Also, just heard from research on the Cohere mappings:

The architecture is similar but there is a very fundamental difference. The MLP block is executed in parallel to the attention.
So the tensor goes through input_layernorm and then from there it goes directly to the attention module and to the MLP module.

I'll create a PR to add the custom mapping for cohere models

[phaelon74]

Awesome, thanks @brian-dellabetta , I was installing from pip install, but I will grab from your main instead.

Thanks also about the cohere models being added. Any ETA on that per se?

[brian-dellabetta]

Hi @phaelon74 , see #1570, you can potentially try on that branch. It looks like it will take ~2hours on a single H100, requires 74GB VRAM for 256 calibration samples with sequence length 512.

[phaelon74]

Hi @phaelon74 , see #1570, you can potentially try on that branch. It looks like it will take ~2hours on a single H100, requires 74GB VRAM for 256 calibration samples with sequence length 512.

Very cool. I wonder then how that extrapolates at 32B and 111B. If an 8B took 74GBVRAM, a 111B ~1024GBVRAM. I'll mess with the 8B model first and then look into how it plays at the bigger models. Thanks again @brian-dellabetta

[brian-dellabetta]

Hi @phaelon74 , the 74GB peak VRAM is for model CohereLabs/c4ai-command-r-plus, which is 104B parameters. With Qwen/Qwen3-30B-A3B it was ~15GB peak VRAM, and most models under that should be similar, it depends on how long vs how wide the LLM is, i.e. lots of decoder layers with small weight matrix ranks will require less memory that fewer decoder layers with high weight matrix ranks.

That's pushing it right to about the threshold of what you can do with a single H100 with our current setup. We are looking into how to improve parallelism, see thread here -- #1382 (comment)