Export HuggingFace mamba to ONNX

[AyoubMDL]

🐛 Describe the bug

Hi, I am trying to convert the Mamba model from Hugging Face to ONNX. However, I encountered the following error:

<class 'RuntimeError'>: Found <class 'transformers.cache_utils.MambaCache'> in output, which is not a known type.

Here is the code:

from transformers import MambaForCausalLM
import torch

model = MambaForCausalLM.from_pretrained("state-spaces/mamba-130m-hf")
dummy_inputs = torch.tensor([[33]], dtype=torch.int32)

with torch.no_grad():
    onnx_program = torch.onnx.export(model,
                                     dummy_inputs,
                                     "mamba_hf.onnx",
                                     input_names=["input"],
                                     opset_version=15,
                                     external_data=False,
                                     dynamo=True)

    onnx_program.optimize()
    onnx_program.save("mamba_hf.onnx")

When I set use_cache=False, the export works without issues. However, the operations differ between the cached (use_cache=True) and non-cached (use_cache=False) cases—specifically, not caching involves different operations (conv1D).

This code works:

from transformers import MambaForCausalLM
import torch

# Disable caching
model = MambaForCausalLM.from_pretrained("state-spaces/mamba-130m-hf", use_cache=False)
dummy_inputs = torch.tensor([[33]], dtype=torch.int32)

with torch.no_grad():
    onnx_program = torch.onnx.export(model,
                                     dummy_inputs,
                                     "mamba_hf.onnx",
                                     input_names=["input"],
                                     opset_version=15,
                                     external_data=False,
                                     dynamo=True)

    onnx_program.optimize()
    onnx_program.save("mamba_hf.onnx")

I am particularly interested in exporting the model with the operations that occur when the cache is available.

Question:

How can I modify the export to ignore the MambaCache object but still retain the logic from this part of the code (with caching):
https://github.com/huggingface/transformers/blob/924f1c717a72261a4b9286a31f199d9512448dd0/src/transformers/models/mamba/modeling_mamba.py#L248

and not this part (without caching):
https://github.com/huggingface/transformers/blob/924f1c717a72261a4b9286a31f199d9512448dd0/src/transformers/models/mamba/modeling_mamba.py#L268

Any guidance on how to handle this would be greatly appreciated!

Versions

PyTorch version: 2.6.0+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: version 3.31.3
Libc version: glibc-2.31

Python version: 3.11.11 (main, Dec 4 2024, 08:55:08) [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-130-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 2080 Ti
Nvidia driver version: 550.107.02
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Byte Order: Little Endian
Address sizes: 39 bits physical, 48 bits virtual
CPU(s): 16
On-line CPU(s) list: 0-15
Thread(s) per core: 2
Core(s) per socket: 8
Socket(s): 1
NUMA node(s): 1
Vendor ID: GenuineIntel
CPU family: 6
Model: 165
Model name: Intel(R) Core(TM) i7-10700 CPU @ 2.90GHz
Stepping: 5
CPU MHz: 4410.645
CPU max MHz: 4800,0000
CPU min MHz: 800,0000
BogoMIPS: 5799.77
Virtualization: VT-x
L1d cache: 256 KiB
L1i cache: 256 KiB
L2 cache: 2 MiB
L3 cache: 16 MiB
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Mitigation; Microcode
Vulnerability Itlb multihit: KVM: Mitigation: VMX disabled
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Mitigation; Enhanced IBRS
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2: Mitigation; Enhanced / Automatic IBRS; IBPB conditional; RSB filling; PBRSB-eIBRS SW sequence; BHI SW loop, KVM SW loop
Vulnerability Srbds: Mitigation; Microcode
Vulnerability Tsx async abort: Not affected
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx rdseed adx smap clflushopt intel_pt xsaveopt xsavec xgetbv1 xsaves dtherm ida arat pln pts hwp hwp_notify hwp_act_window hwp_epp pku ospke md_clear flush_l1d arch_capabilities

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] nvidia-cublas-cu12==12.4.5.8
[pip3] nvidia-cuda-cupti-cu12==12.4.127
[pip3] nvidia-cuda-nvrtc-cu12==12.4.127
[pip3] nvidia-cuda-runtime-cu12==12.4.127
[pip3] nvidia-cudnn-cu12==9.1.0.70
[pip3] nvidia-cufft-cu12==11.2.1.3
[pip3] nvidia-curand-cu12==10.3.5.147
[pip3] nvidia-cusolver-cu12==11.6.1.9
[pip3] nvidia-cusparse-cu12==12.3.1.170
[pip3] nvidia-cusparselt-cu12==0.6.2
[pip3] nvidia-nccl-cu12==2.21.5
[pip3] nvidia-nvjitlink-cu12==12.4.127
[pip3] nvidia-nvtx-cu12==12.4.127
[pip3] onnx==1.16.1
[pip3] onnxconverter-common==1.16.0
[pip3] onnxmltools==1.13.0
[pip3] onnxruntime==1.19.0
[pip3] onnxruntime_extensions==0.13.0
[pip3] onnxscript==0.1.0.dev20250114
[pip3] pytorch-lightning==1.9.5
[pip3] torch==2.6.0
[pip3] torchao==0.8.0
[pip3] torchaudio==2.5.1
[pip3] torchmetrics==1.6.1
[pip3] torchtune==0.5.0
[pip3] torchvision==0.21.0
[pip3] triton==3.2.0
[conda] Could not collect

[justinchuby]

There is a way to override the Cache object. Maybe @titaiwangms knows more?

[titaiwangms]

It's an error from pytree that it does not recognize this object. Basically, you will have to register MambaCache to pytree. Something like this:

def _register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
    # Cache serialization: to be moved into appropriate packages
    import torch

    try:
        from transformers.cache_utils import DynamicCache
    except ImportError:
        DynamicCache = None

    try:
        from transformers.cache_utils import MambaCache
    except ImportError:
        MambaCache = None

    # MambaCache
    unregistered_mamba_cache = True
    if MambaCache is not None and MambaCache in torch.utils._pytree.SUPPORTED_NODES:
        # It is already registered because bypass_export_some_errors was called
        # within a section already calling bypass_export_some_errors or transformers
        # has updated its code to do it.
        # No need to register and unregister then.
        unregistered_mamba_cache = False
    else:
        if verbose:
            print("[bypass_export_some_errors] register MambaCache")
        torch.utils._pytree.register_pytree_node(
            MambaCache,
            flatten_mamba_cache,
            unflatten_mamba_cache,
            serialized_type_name=f"{MambaCache.__module__}.{MambaCache.__name__}",
            flatten_with_keys_fn=flatten_with_keys_mamba_cache,
        )

    # DynamicCache
    unregistered_dynamic_cache = True
    if DynamicCache is not None and DynamicCache in torch.utils._pytree.SUPPORTED_NODES:
        unregistered_dynamic_cache = False
    else:
        from bash_bench.patches.patch_transformers import patched_DynamicCache

        if verbose:
            print("[bypass_export_some_errors] register DynamicCache")
        torch.utils._pytree.register_pytree_node(
            DynamicCache,
            flatten_dynamic_cache,
            unflatten_dynamic_cache,
            serialized_type_name=f"{DynamicCache.__module__}.{DynamicCache.__name__}",
            flatten_with_keys_fn=flatten_with_keys_dynamic_cache,
        )
        torch.fx._pytree.register_pytree_flatten_spec(
            DynamicCache, lambda x, _: [x.key_cache, x.value_cache]
        )

        if verbose:
            print("[bypass_export_some_errors] register patched_DynamicCache")

        torch.utils._pytree.register_pytree_node(
            patched_DynamicCache,
            flatten_dynamic_cache,
            unflatten_pached_dynamic_cache,
            serialized_type_name=f"{patched_DynamicCache.__module__}.{patched_DynamicCache.__name__}",
            flatten_with_keys_fn=flatten_with_keys_dynamic_cache,
        )
        torch.fx._pytree.register_pytree_flatten_spec(
            patched_DynamicCache, lambda x, _: [x.key_cache, x.value_cache]
        )

    return dict(DynamicCache=unregistered_dynamic_cache, MambaCache=unregistered_mamba_cache)

NOTE: This error should be generated at torch.export.export stage

cc @xadupre maybe we should discuss with export team about cache support as well.

[AyoubMDL]

Thanks to @justinchuby and @titaiwangms for the response. Could you please explain how to define the flatten/unflatten functions, or direct me to the relevant documentation ?

[AyoubMDL]

For info, here is an example of flatten/unflatten functions

def flatten_mamba_cache(mamba_cache: MambaCache):
    return ([mamba_cache.intermediate_size,
             mamba_cache.ssm_state_size,
             mamba_cache.conv_kernel_size,
             mamba_cache.batch_size,
             mamba_cache.max_batch_size], None)


def unflatten_mamba_cache(values, to_ignore):
    kwargs = {"intermediate_size": values[0],
              "ssm_state_size": values[1],
              "conv_kernel_size": values[2],
              }
    mamba_config = MambaConfig(**kwargs)
    cache = MambaCache(config=mamba_config,
                       batch_size=values[3],
                       max_batch_size=values[4])
    return cache


def flatten_mamba_with_keys_fn(mamba_cache: MambaCache):
    values, context = flatten_mamba_cache(mamba_cache)
    keys = ['intermediate_size', 'ssm_state_size',
            'conv_kernel_size', 'batch_size', 'max_batch_size']
    key_children = [(torch.utils._pytree.MappingKey(key), value)
                    for key, value in zip(keys, values)]
    return key_children, context

[AyoubMDL]

The ONNX export failed at the third step (Translate to ONNX). Here is the code :

def flatten_mamba_cache(mamba_cache: MambaCache):
    return ([mamba_cache.intermediate_size,
             mamba_cache.ssm_state_size,
             mamba_cache.conv_kernel_size,
             mamba_cache.batch_size,
             mamba_cache.max_batch_size], None)


def unflatten_mamba_cache(values, to_ignore):
    kwargs = {"intermediate_size": values[0],
              "ssm_state_size": values[1],
              "conv_kernel_size": values[2],
              }
    mamba_config = MambaConfig(**kwargs)
    cache = MambaCache(config=mamba_config,
                       batch_size=values[3],
                       max_batch_size=values[4])
    return cache


def flatten_mamba_with_keys_fn(mamba_cache: MambaCache):
    values, context = flatten_mamba_cache(mamba_cache)
    keys = ['intermediate_size', 'ssm_state_size',
            'conv_kernel_size', 'batch_size', 'max_batch_size']
    key_children = [(torch.utils._pytree.MappingKey(key), value)
                    for key, value in zip(keys, values)]
    return key_children, context


def _register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
    # Cache serialization: to be moved into appropriate packages
    import torch

    try:
        from transformers.cache_utils import MambaCache
    except ImportError:
        MambaCache = None

    # MambaCache
    unregistered_mamba_cache = True
    if MambaCache is not None and MambaCache in torch.utils._pytree.SUPPORTED_NODES:
        # It is already registered because bypass_export_some_errors was called
        # within a section already calling bypass_export_some_errors or transformers
        # has updated its code to do it.
        # No need to register and unregister then.
        unregistered_mamba_cache = False
    else:
        if verbose:
            print("[bypass_export_some_errors] register MambaCache")
        torch.utils._pytree.register_pytree_node(
            MambaCache,
            flatten_mamba_cache,
            unflatten_mamba_cache,
            serialized_type_name=f"{MambaCache.__module__}.{MambaCache.__name__}",
            flatten_with_keys_fn=flatten_mamba_with_keys_fn,
        )

    return dict(MambaCache=unregistered_mamba_cache)


_register_cache_serialization()

model = MambaForCausalLM.from_pretrained("state-spaces/mamba-130m-hf")
mamba_cache = MambaCache(model.config, max_batch_size=1)
dummy_inputs = torch.tensor([[33]], dtype=torch.int32)
cache_position = torch.arange(0, 0, device="cpu")

with torch.no_grad():
    onnx_program = torch.onnx.export(model,
                                     (dummy_inputs, None,
                                      None, mamba_cache, None,
                                      None, None, None, cache_position),
                                     "mamba_hf.onnx",
                                     input_names=["input"],
                                     opset_version=15,
                                     external_data=False,
                                     dynamo=True,
                                     report=True)

    onnx_program.optimize()
    onnx_program.save("mamba_hf.onnx")

Here is the report

onnx_export_2025-02-11_17-05-32-540314_conversion.md .

Can you please check what is wrong ? (FYI, the torch forward with the same inputs is working).

[titaiwangms]

Try this

from typing import Any, Dict, List, Tuple

import torch
import transformers
from transformers.cache_utils import MambaCache


def flatten_mamba_cache(
    mamba_cache: "MambaCache",  # noqa: F821
) -> Tuple[List[Any], "torch.utils._pytree.Context"]:
    """Serializes a :class:`transformers.cache_utils.MambaCache` with python objects."""
    flat = [
        (k, getattr(mamba_cache, k))
        for k in [
            # "max_batch_size",  # new in transformers==4.47
            # "intermediate_size",
            # "ssm_state_size",
            # "conv_kernel_size",
            "conv_states",
            "ssm_states",
        ]
        if hasattr(mamba_cache, k)
    ]
    return [f[1] for f in flat], [f[0] for f in flat]


def unflatten_mamba_cache(
    values: List[Any],
    context: "torch.utils._pytree.Context",
    output_type=None,
) -> "MambaCache":  # noqa: F821
    """Restores a :class:`transformers.cache_utils.MambaCache` from python objects."""
    conv_states, ssm_states = values

    class _config:  # noqa: N801
        def __init__(self):
            self.intermediate_size = conv_states.shape[2]
            self.state_size = ssm_states.shape[3]
            self.conv_kernel = conv_states.shape[3]
            self.num_hidden_layers = conv_states.shape[0]

    from transformers.cache_utils import MambaCache

    cache = MambaCache(_config(), batch_size=1, dtype=values[-1].dtype)
    values = dict(zip(context, values))
    for k, v in values.items():
        setattr(cache, k, v)
    return cache


def flatten_mamba_with_keys_fn(mamba_cache: MambaCache):
    values, context = flatten_mamba_cache(mamba_cache)
    keys = ['intermediate_size', 'ssm_state_size',
            'conv_kernel_size', 'batch_size', 'max_batch_size']
    key_children = [(torch.utils._pytree.MappingKey(key), value)
                    for key, value in zip(keys, values)]
    return key_children, context


def _register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
    # Cache serialization: to be moved into appropriate packages
    import torch

    try:
        from transformers.cache_utils import MambaCache
    except ImportError:
        MambaCache = None

    # MambaCache
    unregistered_mamba_cache = True
    if MambaCache is not None and MambaCache in torch.utils._pytree.SUPPORTED_NODES:
        # It is already registered because bypass_export_some_errors was called
        # within a section already calling bypass_export_some_errors or transformers
        # has updated its code to do it.
        # No need to register and unregister then.
        unregistered_mamba_cache = False
    else:
        if verbose:
            print("[bypass_export_some_errors] register MambaCache")
        torch.utils._pytree.register_pytree_node(
            MambaCache,
            flatten_mamba_cache,
            unflatten_mamba_cache,
            serialized_type_name=f"{MambaCache.__module__}.{MambaCache.__name__}",
            flatten_with_keys_fn=flatten_mamba_with_keys_fn,
        )

    return dict(MambaCache=unregistered_mamba_cache)


_register_cache_serialization()

model = transformers.MambaForCausalLM.from_pretrained("state-spaces/mamba-130m-hf")
mamba_cache = MambaCache(model.config, max_batch_size=1)
dummy_inputs = torch.tensor([[33]], dtype=torch.int32)
cache_position = torch.arange(0, 0, device="cpu")

with torch.no_grad():
    onnx_program = torch.onnx.export(model,
                                     (dummy_inputs, None,
                                      None, mamba_cache, None,
                                      None, None, None, cache_position),
                                     "mamba_hf.onnx",
                                     input_names=["input"],
                                     opset_version=15,
                                     external_data=False,
                                     dynamo=True,
                                     optimize=True,
                                     report=True)

Tagging the patch author @xadupre to see if there is anything to add.

[AyoubMDL]

Thanks @titaiwangms , I still get the same error. I think it is related to None values (the line that throw the error https://github.com/microsoft/onnxscript/blob/71e24d9b087be50f5aefc7a5b66bac4298317f42/onnxscript/ir/_convenience.py#L372 as it gets a None value with a defined dtype). Also, there is a node SequenceConstruct that contains some None values as inputs {'inputs': (None, None, SymbolicTensor('clamp', type=Tensor(INT64), shape=[1], producer=node_Min_34, index=0))} I'll be debugging that

[titaiwangms]

Can you try update your onnxscript?

[justinchuby]

Thanks @titaiwangms , I still get the same error. I think it is related to None values (the line that throw the error https://github.com/microsoft/onnxscript/blob/71e24d9b087be50f5aefc7a5b66bac4298317f42/onnxscript/ir/_convenience.py#L372 as it gets a None value with a defined dtype). Also, there is a node SequenceConstruct that contains some None values as inputs {'inputs': (None, None, SymbolicTensor('clamp', type=Tensor(INT64), shape=[1], producer=node_Min_34, index=0))} I'll be debugging that

Please set report=True and attach an error report. Thanks

[AyoubMDL]

Thanks @titaiwangms , @justinchuby . Upgrading onnxscript allows the export to finish. However I ended up with a non-functional model.
Here is the report when I add verify=True

onnx_export_2025-02-12_17-42-10-177935_runtime.md

One observation: the report shows that onnx checker works and the onnxruntime fails, but when I run the onnx checker directly on the model, it fails, here's the error:

C.check_model(
onnx.onnx_cpp2py_export.checker.ValidationError: Node (node_Unsqueeze_64)'s input 0 is marked single but has an empty string in the graph

==> Context: Bad node spec for node. Name: node_Unsqueeze_64 OpType: Unsqueeze

there are some Unsqueeze nodes that have empty data.

[justinchuby]

Thanks. This seems like a bug. Could you share the onnx model (no weights needed) if possible?

[AyoubMDL]

Here is the model

mamba_hf.zip