[Bug]: SequentialPipeline fails on ERNIE-4.5-VL (remote code) with FX trace TypeError: to(device=MetaDeviceAttribute)

[Firworksyt]

⚙️ Your current environment

The output of python collect_env.py

### Environment Information ###
Operating System: `Linux-6.8.0-64-generic-x86_64-with-glibc2.35`
Python Version: `3.11.14 (main, Oct 31 2025, 23:04:14) [Clang 21.1.4 ]`
llm-compressor Version: `0.8.1`
compressed-tensors Version: `0.12.2`
transformers Version: `4.56.2`
torch Version: `2.8.0`
CUDA Devices: `['NVIDIA RTX PRO 6000 Blackwell Server Edition', 'NVIDIA RTX PRO 6000 Blackwell Server Edition']`
AMD Devices: `None`

🐛 Describe the bug

When quantizing ERNIE-4.5-VL-28B-A3B-Thinking (baidu/ERNIE-4.5-VL-28B-A3B-Thinking) with LLM-Compressor, the SequentialPipeline inferred by default crashes during FX tracing with:
TypeError: to() received an invalid combination of arguments - got (device=MetaDeviceAttribute, )

This happens before calibration even starts, on a very small dataset slice.

If I instead force pipeline="basic", oneshot runs (until it eventually hits CUDA OOM for larger calibration settings). So the model itself can be quantized, but the sequential/onloading pipeline is currently incompatible.

Maybe I'm just doing something stupid here but this is generally how I've been quantizing other models to NVFP4.

🛠️ Steps to reproduce

import torch
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM
from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier

MODEL_ID = "baidu/ERNIE-4.5-VL-28B-A3B-Thinking"

# Load tokenizer
tok = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)

# Load model with bf16 dtype to avoid unnecessary FP32 tensors
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    trust_remote_code=True,
)

# Tiny dataset slice to make repro cheap
raw = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft[:2]")

def to_text(ex):
    return {"text": tok.apply_chat_template(ex["messages"], tokenize=False)}

raw = raw.map(to_text)

ds = raw.map(
    lambda s: tok(s["text"], truncation=True, max_length=256),
    remove_columns=raw.column_names,
)

# Minimal NVFP4 recipe
recipe = QuantizationModifier(
    targets="Linear",
    scheme="NVFP4",
    ignore=["lm_head"],
)

# Run oneshot quantization with default (inferred) pipeline
oneshot(
    model=model,
    dataset=ds,
    recipe=recipe,
    num_calibration_samples=2,
    max_seq_length=256,
    trust_remote_code_model=True,
)

Running this should produce the same error. The number of samples and seq length are extremely low just for troubleshooting speed.

[kylesayrs]

Hi there!

This is very likely due to the model not supporting device offloading in its current definition.

The follow lines are examples of where the model attempts to get the device of an offloaded weight, thus resulting in a meta tensor. Patching location like these (and using the activation's device, rather than the weight's device) will fix the issue you're seeing. You can also wait for the HF team or others to do this.
https://huggingface.co/baidu/ERNIE-4.5-VL-28B-A3B-Thinking/blob/main/modeling_ernie4_5_vl.py#L292
https://huggingface.co/baidu/ERNIE-4.5-VL-28B-A3B-Thinking/blob/main/modeling_ernie4_5_vl.py#L726
https://huggingface.co/baidu/ERNIE-4.5-VL-28B-A3B-Thinking/blob/main/modeling_ernie4_5_vl.py#L1206

There are some designs which move away from accelerate's cpu offloading (which will always have these issues). However these designs aren't currently scheduled to implementation by the RedHat model optimization team atm.
https://expensive-date-2f9.notion.site/Replacing-Accelerate-Offloading-25a213393e8580c08e79fbcd215a1c09?pvs=74

[Firworksyt]

Thanks for the quick response!

I implemented the suggested device fix so the MLPs now use the activation device (x.device) instead of weight.device (which was meta under offloading). That completely resolved the earlier MetaDevice .to() crash.

However, after that fix, the quantization run now fails inside the SequentialPipeline AST wrapper:

IndependentPipeline | INFO - Inferred `SequentialPipeline` for `QuantizationModifier`
Traceback (most recent call last):
  File "/home/ubuntu/quantize_nvfp4.py", line 189, in <module>
    oneshot(
  File ".../llmcompressor/entrypoints/oneshot.py", line 330, in oneshot
    one_shot()
  File ".../llmcompressor/entrypoints/oneshot.py", line 158, in __call__
    self.apply_recipe_modifiers(
  File ".../llmcompressor/entrypoints/oneshot.py", line 201, in apply_recipe_modifiers
    pipeline(
  File ".../llmcompressor/pipelines/independent/pipeline.py", line 45, in __call__
    pipeline(model, dataloader, dataset_args)
  File ".../llmcompressor/pipelines/sequential/pipeline.py", line 71, in __call__
    subgraphs = trace_subgraphs(model, sample_input, sequential_targets, ignore)
  File ".../llmcompressor/pipelines/sequential/helpers.py", line 120, in trace_subgraphs
    stack.enter_context(autowrap_forwards(ancestors, ignore))
  File ".../llmcompressor/pipelines/sequential/ast_helpers.py", line 30, in autowrap_forwards
    stack.enter_context(autowrap_forward(module, ignore))
  File ".../llmcompressor/pipelines/sequential/ast_helpers.py", line 77, in autowrap_forward
    new_forward = namespace["forward"].__get__(module)
KeyError: 'forward'

This new failure happens during the AST rewrite phase (autowrap_forward) when SequentialPipeline attempts to wrap module forward methods.

ERNIE-4.5-VL requires a couple of small patches to run normally:
A RoPE3D position_ids patch
A patch to the base model .forward to handle 3D+ inputs_embeds
After your suggested fix, I now also apply an MLP offload-device patch (activation device instead of weight device)

I put together a minimal standalone repro script that hits the same KeyError: 'forward' in autowrap_forward using ERNIE-4.5-VL with the three patches.

import os
import torch
import torch.nn.functional as F
from types import MethodType

from datasets import Dataset
from transformers import AutoTokenizer, AutoModelForCausalLM
from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier

MODEL_ID = "baidu/ERNIE-4.5-VL-28B-A3B-Thinking"

# Keep Inductor/Dynamo out of the way
os.environ.setdefault("TORCHINDUCTOR_DISABLE", "1")
os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")


# ---- Patch 1: RoPE3D position_ids for pure-text inputs ----
def patch_ernie_vl_rope3d(model):
    if model.__class__.__name__ != "Ernie4_5_VLMoeForConditionalGeneration":
        print("[INFO] Not an ERNIE-4.5-VL model; skipping RoPE3D patch.")
        return model

    print("[INFO] Applying RoPE3D position_ids patch for ERNIE-4.5-VL...")

    orig_forward = model.forward

    def _patched_forward(self, *args, **kwargs):
        input_ids = kwargs.get("input_ids", None)

        if input_ids is not None and isinstance(input_ids, torch.Tensor):
            if input_ids.dim() != 2:
                bsz = input_ids.shape[0]
                seq_len = input_ids.numel() // bsz
                input_ids = input_ids.view(bsz, seq_len)
                kwargs["input_ids"] = input_ids

            if getattr(self.config, "rope_3d", False):
                pos = kwargs.get("position_ids", None)
                if pos is None:
                    bsz, seq_len = input_ids.shape
                    device = input_ids.device

                    t = torch.arange(seq_len, device=device).unsqueeze(0).expand(bsz, -1)
                    h = torch.zeros_like(t)
                    w = torch.zeros_like(t)

                    position_ids = torch.stack([t, h, w], dim=-1)  # [bsz, seq, 3]
                    kwargs["position_ids"] = position_ids

        return orig_forward(*args, **kwargs)

    model.forward = MethodType(_patched_forward, model)
    return model


# ---- Patch 2: base_model.forward for >3D inputs_embeds ----
def patch_ernie_vl_inputs_embeds(model):
    base_model = getattr(model, "model", None)
    if base_model is None:
        base_model = model

    _orig_forward = base_model.forward

    def _patched_forward(self, *args, **kwargs):
        inputs_embeds = kwargs.get("inputs_embeds", None)

        if inputs_embeds is not None and isinstance(inputs_embeds, torch.Tensor):
            if inputs_embeds.dim() > 3:
                bsz = inputs_embeds.shape[0]
                hidden = inputs_embeds.shape[-1]
                seq = inputs_embeds.numel() // (bsz * hidden)
                inputs_embeds = inputs_embeds.view(bsz, seq, hidden)
                kwargs["inputs_embeds"] = inputs_embeds

        return _orig_forward(*args, **kwargs)

    base_model.forward = MethodType(_patched_forward, base_model)
    print("[INFO] Patched Ernie4_5 base model.forward to handle >3D inputs_embeds")
    return model


# ---- Patch 3: MLP device patch (weight.device -> x.device) ----
def patch_ernie_vl_offload_devices(model):
    """
    Fix ERNIE 4.5 VL MLP offloading bugs by replacing references to
    weight.device (which may be meta) with activation device (x.device).
    """

    def patch_mlp_module(mlp_module):
        orig_class = mlp_module.__class__.__name__

        def _forward(self, x):
            # Use activation's device, don't depend on weight.device (meta under offload)
            x = x.to(x.device)

            if orig_class == "Ernie4_5_MLP":
                return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))

            # MoE MLP
            x_local = F.silu(self.gate_proj(x)) * self.up_proj(x)

            if getattr(self, "moe_dropout_prob", 0) > 0:
                x_local = F.dropout(
                    input=x_local,
                    p=self.moe_dropout_prob,
                    training=self.training,
                )

            return self.down_proj(x_local)

        mlp_module.forward = MethodType(_forward, mlp_module)

    patched = 0
    for module in model.modules():
        cls = module.__class__.__name__
        if cls in ("Ernie4_5_MLP", "Ernie4_5_MoeMLP"):
            patch_mlp_module(module)
            patched += 1

    print(f"[INFO] Patched {patched} Ernie4_5_*MLP modules for offload device safety")
    return model


def main():
    print(f"[INFO] Loading tokenizer and model: {MODEL_ID}")
    tok = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)

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

    model.eval()
    torch.set_grad_enabled(False)

    # Apply patches
    model = patch_ernie_vl_rope3d(model)
    model = patch_ernie_vl_inputs_embeds(model)
    model = patch_ernie_vl_offload_devices(model)

    # Tiny synthetic dataset (no HF download)
    texts = [
        "Hello, this is a short test for ERNIE VL.",
        "Another calibration sample for NVFP4 sequential pipeline.",
    ]
    enc = tok(
        texts,
        truncation=True,
        max_length=256,
        padding=False,
    )

    ds = Dataset.from_dict(
        {
            "input_ids": enc["input_ids"],
            "attention_mask": enc["attention_mask"],
        }
    )

    print("[INFO] Building NVFP4 QuantizationModifier")
    recipe = QuantizationModifier(
        targets="Linear",
        scheme="NVFP4",
        ignore=["lm_head"],
    )

    print("[INFO] Calling oneshot() with inferred SequentialPipeline")
    oneshot(
        model=model,
        dataset=ds,
        recipe=recipe,
        num_calibration_samples=2,
        max_seq_length=256,
        trust_remote_code_model=True,
        # NOTE: no pipeline="basic" here, we want Sequential
    )


if __name__ == "__main__":
    main()

Running that script should reproduce the same KeyError: 'forward' failure inside the SequentialPipeline AST wrapper.

[kylesayrs]

Seem like there are some modules which do not implement a forward method? This is similar to an issue which occurred with the qwen thinker model as well.

Best solution is to patch the AutoWrapper logic to conditionally check for a forward method.

[Firworksyt]

After guarding for missing forward in ast_helper.py and applying the other patches, SequentialPipeline now gets to FX tracing but fails with:

Traceback (most recent call last):
  File "/workspace/quantize_nvfp4.py", line 143, in <module>
    oneshot(
  File "/workspace/.venv/lib/python3.11/site-packages/llmcompressor/entrypoints/oneshot.py", line 330, in oneshot
    one_shot()
  File "/workspace/.venv/lib/python3.11/site-packages/llmcompressor/entrypoints/oneshot.py", line 158, in __call__
    self.apply_recipe_modifiers(
  File "/workspace/.venv/lib/python3.11/site-packages/llmcompressor/entrypoints/oneshot.py", line 201, in apply_recipe_modifiers   
    pipeline(
  File "/workspace/.venv/lib/python3.11/site-packages/llmcompressor/pipelines/independent/pipeline.py", line 45, in __call__       
    pipeline(model, dataloader, dataset_args)
  File "/workspace/.venv/lib/python3.11/site-packages/llmcompressor/pipelines/sequential/pipeline.py", line 71, in __call__        
    subgraphs = trace_subgraphs(model, sample_input, sequential_targets, ignore)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/.venv/lib/python3.11/site-packages/llmcompressor/pipelines/sequential/helpers.py", line 125, in trace_subgraphs 
    tracer.trace(
  File "/workspace/.venv/lib/python3.11/site-packages/transformers/utils/fx.py", line 1314, in trace
    self.graph = super().trace(root, concrete_args=concrete_args)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/.venv/lib/python3.11/site-packages/torch/fx/_symbolic_trace.py", line 850, in trace
    (self.create_arg(fn(*args)),),
                     ^^^^^^^^^
  File "/workspace/quantize_nvfp4.py", line -1, in _patched_forward
  File "/workspace/.venv/lib/python3.11/site-packages/torch/fx/proxy.py", line 542, in __iter__
    return self.tracer.iter(self)
           ^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/.venv/lib/python3.11/site-packages/torch/fx/proxy.py", line 399, in iter
    raise TraceError(
torch.fx.proxy.TraceError: Proxy object cannot be iterated. This can be attempted when the Proxy is used in a loop or as a *args or **kwargs function argument. See the torch.fx docs on pytorch.org for a more detailed explanation of what types of control flow can be traced, and check out the Proxy docstring for help troubleshooting Proxy iteration errors

I'm not sure if there's anything else I could try to be able to quantize this model or if it should just be considered incompatible with the SequentialPipeline in LLM-Compressor at this time.
If you have any suggestions for workarounds (e.g., specific modules to exclude from tracing, or patterns known to cause Proxy iteration issues), I'm happy to try them.

[kylesayrs]

This happens when the model is doing some sort of iteration, which cannot be traced. You can put a breakpoint in torch/fx/proxy.py then use where to figure out which line is causing the issue.

If you run with these changes, you'll be able to see the stack trace of the autowrapped code more clearly.