Saving model is very slow when setting save_compressed=True

[wangwenmingaa]

After quantizing the model using dynamic FP8 quantization, I found that saving the model with save_compressed=True is extremely slow. Why is this slow? What operations are performed during this process? What is the difference compared to setting it to False? Is there a way to improve the saving speed?

[wangwenmingaa]

What is the difference between a quantized model saved with save_compressed=False and the original model?

[wangwenmingaa]

I spent over 4 hours, but only completed 10,233 iterations. The code I ran is as follows:
Compressing model: 10233it [4:21:11, 1.54it/s]

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

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

MODEL_ID = "/open_source/Qwen3-30B-A3B"

# Load model.
# model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto",device_map="auto")
# tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

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

# Configure the quantization algorithm and scheme.
# In this case, we:
#   * quantize the weights to fp4 with per group 16 via ptq
recipe = QuantizationModifier(targets="Linear", scheme="NVFP4A16", ignore=["lm_head", "re:.*mlp.gate$", "re:.*mlp.shared_expert_gate$"],)


output_dir = "./Qwen3-30B-A3B-NVFP4-W4A16-0819"
# Apply quantization.
oneshot(
    model=model,
    recipe=recipe,
    save_compressed=True,
    trust_remote_code_model=True,
    output_dir=output_dir,
)

[dsikka]

Hi @wangwenmingaa - are you using a GPU for your script?

When you save the model, the compressor is applied and the model is saved in the nvfp4-pack-quantized format. This compresses the model and makes it runnable in vllm where it can leverage performant kernels for quantization.

If you set save_compressed as False, the model will be saved with its scales and zero-points but its weights will remain dense and the compressor will not be applied. Setting save_compressed to True applies the compressor to compress the weights

[brian-dellabetta]

Hi @wangwenmingaa , can you try on this branch of compressed-tensors? neuralmagic/compressed-tensors#400

This adds torch.compile to speed up the packing for nvfp4, but could use further validation

[wangwenmingaa]

Hi @wangwenmingaa - are you using a GPU for your script?

When you save the model, the compressor is applied and the model is saved in the nvfp4-pack-quantized format. This compresses the model and makes it runnable in vllm where it can leverage performant kernels for quantization.

If you set save_compressed as False, the model will be saved with its scales and zero-points but its weights will remain dense and the compressor will not be applied. Setting save_compressed to True applies the compressor to compress the weights

@dsikka Thank you for your reply. I'm using 2*A100 GPUs, but during the execution of "Compressing model: 10233it [4:21:11, 1.54it/s]", the computation is running on the CPU instead of the GPU. My command is python script_name.py. Could there be an issue with the command I'm running?

[wangwenmingaa]

Hi @wangwenmingaa , can you try on this branch of compressed-tensors? neuralmagic/compressed-tensors#400

This adds torch.compile to speed up the packing for nvfp4, but could use further validation

Alright, I'll try it out.

[dsikka]

@wangwenmingaa Can you also update how you're loading the model, to this:

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, torch_dtype=torch.bfloat16, trust_remote_code=True,
)