Adapting AutoRound for a New LLM Architecture

Overview

Most standard Transformers-based LLMs work with AutoRound out-of-the-box. This skill covers what to do when a new model architecture requires code changes. The need for adaptation typically arises from:

• Non-standard layer hierarchy (block detection fails)
• Fused Mixture-of-Experts (MoE) weights
• Non-standard linear layer types (not nn.Linear or Conv1D)
• Complex multi-component architectures (multimodal routing)
• Shared cache keys or position embeddings

Step 0: Diagnose the Problem

Try quantizing the model first:

from auto_round import AutoRound

ar = AutoRound("your-org/your-model", scheme="W4A16", iters=2, nsamples=2)
ar.quantize_and_save(output_dir="./test_output", format="auto_round")

Common failure modes and their fixes:

Step 1: Fix Block Detection

AutoRound discovers quantizable blocks via get_block_names() which searches recursively for nn.ModuleList instances. If your model has a non-standard layer hierarchy, block detection may fail.

Check current detection

from auto_round.utils import get_block_names

model = ...  # loaded model
print(get_block_names(model))

Option A: Use to_quant_block_names parameter

For simple cases, override block names without code changes:

ar = AutoRound(
    model,
    to_quant_block_names="model.decoder.layers",  # explicit path
)

Option B: Register in SPECIAL_MULTIMODAL_BLOCK

For multimodal or multi-component models, add a custom block handler in auto_round/special_model_handler.py:

def _get_your_model_multimodal_block(model, quant_vision=False):
    """Get block names for YourModel.

    YourModel structure:
    - encoder.layers: encoder blocks
    - decoder.layers: decoder blocks
    """
    block_names = []

    if quant_vision and hasattr(model, "encoder"):
        block_names.append([f"encoder.layers.{i}" for i in range(len(model.encoder.layers))])

    block_names.append([f"decoder.layers.{i}" for i in range(len(model.decoder.layers))])

    return block_names


# Register: key must match model.config.model_type
SPECIAL_MULTIMODAL_BLOCK["your_model_type"] = _get_your_model_multimodal_block

Also add to support lists if applicable:

# If text-only calibration works for this multimodal model:
SUPPORT_ONLY_TEXT_MODELS.append("your_model_type")

# If batch_size must be limited:
mllms_with_limited_bs = (..., "your_model_type")

Step 2: Handle MoE (Mixture-of-Experts) Models

MoE models often have fused 3D expert weights (shape [num_experts, hidden, intermediate]) that must be "unfused" into per-expert nn.Linear layers for quantization.

Check if auto-handled

Transformers >= 5.0 has a linear_loop experts interface that auto-unfuses most MoE models. Test first — it may just work.

Register custom unfusing

If auto-unfusing fails, create a custom module in auto_round/modeling/fused_moe/:

1. Create auto_round/modeling/fused_moe/your_moe.py:

"""Unfuse fused MoE weights for YourModel."""

import torch
import torch.nn as nn
from auto_round.modeling.fused_moe.replace_modules import register_replacement


@register_replacement("YourMoELayer")
def replace_your_moe_layer(module, name, model):
    """Replace FusedMoE with per-expert nn.Linear layers."""
    experts = nn.ModuleList()
    for i in range(module.num_experts):
        linear = nn.Linear(module.hidden_size, module.intermediate_size, bias=False)
        linear.weight.data = module.weight[i].clone()
        experts.append(linear)
    return experts

2. Register in BUILTIN_MODULES:

Edit auto_round/modeling/fused_moe/replace_modules.py:

BUILTIN_MODULES["your_model_type"] = LazyImport("auto_round.modeling.fused_moe.your_moe")

Existing MoE implementations

Step 3: Add Custom Forward Pass

Some models have non-standard forward passes that don't get calibrated correctly with the default model.forward(). This is common for multi-component architectures.

Edit _handle_special_model() in auto_round/special_model_handler.py:

def _your_model_forward(model, **kwargs):
    """Custom forward that routes through all quantizable components."""
    # Example: route through both encoder and decoder
    encoder_output = model.encoder(**kwargs)
    decoder_output = model.decoder(encoder_output, **kwargs)
    return decoder_output


def _handle_special_model(model):
    ...
    if hasattr(model, "config") and model.config.model_type == "your_model_type":
        from functools import partial

        model.forward = partial(_your_model_forward, model)
    return model

When is this needed?

• Model has multiple sub-models (thinker/talker, encoder/decoder)
• Default forward doesn't exercise all quantizable layers
• Model needs special input preprocessing during calibration

Existing examples

Step 4: Handle Non-Standard Linear Layers

AutoRound quantizes these layer types by default:

# auto_round/utils/common.py
SUPPORTED_LAYER_TYPES = (torch.nn.Linear, transformers.pytorch_utils.Conv1D)
INNER_SUPPORTED_LAYER_TYPES = ("FP8Linear",)  # matched by class name string

If your model uses a custom linear type (e.g., QuantizedLinear, FP8Linear), it won't be quantized unless registered.

Option A: String-based matching

INNER_SUPPORTED_LAYER_TYPES matches by class name string — useful for external classes that can't be imported directly:

INNER_SUPPORTED_LAYER_TYPES = ("FP8Linear", "YourCustomLinear")

Option B: Type-based registration

If you can import the class:

from your_library import YourLinear

SUPPORTED_LAYER_TYPES = SUPPORTED_LAYER_TYPES + (YourLinear,)

Step 5: Handle Shared Cache Keys

Some models share tensors across blocks during inference (e.g., Gemma3's rotary position embeddings). These must be declared so the calibration cache doesn't duplicate or corrupt them.

Edit SPECIAL_SHARED_CACHE_KEYS in auto_round/special_model_handler.py:

SPECIAL_SHARED_CACHE_KEYS["YourModelForCausalLM"] = ("shared_position_embeddings", "shared_rope")

The key is the class name of the model (not model_type).

Step 6: Test

def test_your_model_quantization():
    ar = AutoRound(
        "your-org/your-model",
        scheme="W4A16",
        iters=2,
        nsamples=2,
        batch_size=2,
    )
    compressed_model, layer_config = ar.quantize()
    # Verify layers were quantized
    assert len(layer_config) > 0, "No layers were quantized"

    ar.save_quantized(output_dir="./tmp_your_model", format="auto_round")

    # Verify inference works
    from auto_round.utils import model_infer

    output = model_infer(compressed_model, tokenizer, "Hello world")
    assert output is not None

Step 7: Update Documentation

1. Add model to supported list in README.md
2. Update README_CN.md with equivalent Chinese content
3. Add example script if the model has notable differences

Checklist

• get_block_names() finds all quantizable blocks
• MoE layers (if any) are unfused correctly
• calib() runs without shape errors
• All target layers are quantized (check "Quantized X/Y layers" log)
• Forward pass exercises all quantizable components
• Quantized model produces valid outputs
• Export to target format works
• README.md + README_CN.md updated

Key Files