[Bug]: torch.OutOfMemoryError: CUDA out of memory for quart Qwen3-moe-30B-A3B

[QuanhuiGuan]

🐛 Describe the bug

I used four graphics cards and set the device_map parameter, but only one graphics card is used during the quantization process.

🛠️ Steps to reproduce

import torch
from datasets import load_dataset
from datasets import Dataset
# from random import shuffle
from transformers import AutoModelForCausalLM, AutoTokenizer

from llmcompressor import oneshot
from llmcompressor.modifiers.awq import AWQModifier
from llmcompressor.utils import dispatch_for_generation
from llmcompressor.transformers.compression.helpers import calculate_offload_device_map

# Select model and load it.
MODEL_ID = "/cloud/data5/model_gallery/Qwen3-30B-A3B"



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

# Select calibration dataset.
DATASET_ID = "/cloud/data5/data_gallery/ultrachat_200k"
DATASET_SPLIT = "train_sft"

# 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 = 10280

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


def preprocess(example):
    return {
        "text": tokenizer.apply_chat_template(
            example["messages"],
            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.
# NOTE: vllm currently does not support asym MoE, using symmetric here
recipe = [
    AWQModifier(
        ignore=["lm_head", "re:.*mlp.gate$", "re:.*mlp.shared_expert_gate$"],
        scheme="W4A16",
        targets=["Linear"],
    ),
]

# 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 ==============")
dispatch_for_generation(model)
input_ids = tokenizer("请问，你是谁？", 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")

# Save to disk compressed.
SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-awq-sym"
model.save_pretrained(SAVE_DIR, save_compressed=True)
tokenizer.save_pretrained(SAVE_DIR)

[CAlexander0614]

I also have the same exact issue with a Mistral 3.1 small model and am using a similar MAX_SEQUENCE_LENGTH. When I reduced it to 2048 it ran fine however this is not a feasible solution for me as I need the model to take in 5 1080p images.

import json
import os

import requests
import torch
from PIL import Image
from transformers import AutoProcessor, Mistral3ForConditionalGeneration

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

model_id = "model3"
model = Mistral3ForConditionalGeneration.from_pretrained(model_id, torch_dtype="auto")
processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)

file_path = "mistral3_chat_template.json"
with open(file_path, "r") as file:
    processor.chat_template = json.load(file)["chat_template"]

DATASET_ID = "flickr30k"
DATASET_SPLIT = "test"
NUM_CALIBRATION_SAMPLES = 128
MAX_SEQUENCE_LENGTH = 10068


def data_collator(batch):
    assert len(batch) == 1
    return {
        key: torch.tensor(value)
        if key != "pixel_values"
        else torch.tensor(value, dtype=model.dtype)
        for key, value in batch[0].items()
    }


recipe = [
    GPTQModifier(
        targets="Linear",
        scheme="W8A8",
        ignore=["re:.*lm_head", "re:.*multi_modal_projector.*", "re:.*vision_model.*"]
    ),
]

oneshot(
    model=model,
    tokenizer=model_id,
    dataset=DATASET_ID,
    splits={"calibration": f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES}]"},
    recipe=recipe,
    max_seq_length=MAX_SEQUENCE_LENGTH,
    num_calibration_samples=NUM_CALIBRATION_SAMPLES,
    trust_remote_code_model=True,
    data_collator=data_collator,
    sequential_targets=["MistralDecoderLayer"],
)

print("========== SAMPLE GENERATION ==============")
dispatch_for_generation(model)
messages = [
    {
        "role": "user",
        "content": [
            {"type": "text", "text": "Please describe the animal in this image\n"},
            {"type": "image"},
        ],
    },
]
prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
image_url = "http://images.cocodataset.org/train2017/000000231895.jpg"
raw_image = Image.open(requests.get(image_url, stream=True).raw)

inputs = processor(images=raw_image, text=prompt, return_tensors="pt").to(model.device)
inputs["pixel_values"] = inputs["pixel_values"].to(model.dtype)  # fix dtype
output = model.generate(**inputs, max_new_tokens=100)
print(processor.decode(output[0], skip_special_tokens=True))
print("==========================================")

SAVE_DIR = model_id.rstrip("/").split("/")[-1] + "-W4A16-G128"
model.save_pretrained(SAVE_DIR, save_compressed=True)
processor.save_pretrained(SAVE_DIR)

[brian-dellabetta]

Hi @QuanhuiGuan and @CAlexander0614 , AWQ and GPTQ currently operate sequentially, layer-by-layer through the model. It is best to not set device_map when loading the model, our pipeline willl handle this layerwise onloading by itself. Neither algorithm is parallelized to run on multiple GPUs at present, but you can reduce VRAM requirements by offloading values stored during the normal operation of the algorithms.

• GPTQ offload_hessians=True

  llm-compressor/src/llmcompressor/modifiers/quantization/gptq/base.py

  Line 82 in 3e721cb

  :param offload_hessians: Set to True for decreased memory usage but increased

• AWQ offload_device=torch.device("cpu")

  llm-compressor/src/llmcompressor/modifiers/awq/base.py

  Line 109 in 3e721cb

  :param offload_device: offload cached args to this device, which reduces memory

MoE models will require a lot of caching, try adding these and see if you can use the num_samples/sequence_length you're hoping for.

[CAlexander0614]

Hi @QuanhuiGuan and @CAlexander0614 , AWQ and GPTQ currently operate sequentially, layer-by-layer through the model. It is best to not set device_map when loading the model, our pipeline willl handle this layerwise onloading by itself. Neither algorithm is parallelized to run on multiple GPUs at present, but you can reduce VRAM requirements by offloading values stored during the normal operation of the algorithms.

• GPTQ offload_hessians=True

    [llm-compressor/src/llmcompressor/modifiers/quantization/gptq/base.py](https://github.com/vllm-project/llm-compressor/blob/3e721cb741cf234d0c28ee265664275ab6b03857/src/llmcompressor/modifiers/quantization/gptq/base.py#L82)
  
  
       Line 82
    in
    [3e721cb](/vllm-project/llm-compressor/commit/3e721cb741cf234d0c28ee265664275ab6b03857)
  
  
  
  
  
      
        
             :param offload_hessians: Set to True for decreased memory usage but increased
  

• AWQ offload_device=torch.device("cpu")

    [llm-compressor/src/llmcompressor/modifiers/awq/base.py](https://github.com/vllm-project/llm-compressor/blob/3e721cb741cf234d0c28ee265664275ab6b03857/src/llmcompressor/modifiers/awq/base.py#L109)
  
  
       Line 109
    in
    [3e721cb](/vllm-project/llm-compressor/commit/3e721cb741cf234d0c28ee265664275ab6b03857)
  
  
  
  
  
      
        
             :param offload_device: offload cached args to this device, which reduces memory
  

MoE models will require a lot of caching, try adding these and see if you can use the num_samples/sequence_length you're hoping for.

Thank you offload_hessians=True seems to be working fine