pytorch - 💡(How to fix) Fix run_decomposition fails when using flattening functions [1 participants]

pytorch2026-04-07 09:18:54

ON THIS PAGE

Recommended Tools

×6

Utilities matched from this issue’s tags and category — try them while you read without losing context.

GitHub issue graph ai analysis

Paste a GitHub issue URL. We fetch that issue, discover linked issues from bodies/comments/timeline, collect linked pull requests, and produce a structured English report.

The report is written in English Markdown for sharing and archival.

GitHub issue URL

Helpful · Quick feedback

GitHub stats

pytorch/pytorch#179555•Fetched 2026-04-08 03:00:30

View on GitHub

Comments

Participants

Timeline

Reactions

Author

xadupre

Participants

xadupre

Timeline (top)

mentioned ×9subscribed ×9labeled ×3

Error Message

def test_mistral_nemo_torch_exoprt_export_issue(): import itertools import torch import transformers from transformers import AutoModelForCausalLM, MistralConfig

num_hidden_layers = 1
config = MistralConfig(
    architectures=["MistralNeMoForCausalLM"],
    bos_token_id=1,
    eos_token_id=2,
    hidden_act="silu",
    hidden_size=512,
    intermediate_size=1376,
    max_position_embeddings=1024,
    model_type="mistral",
    num_attention_heads=8,
    num_hidden_layers=num_hidden_layers,
    num_key_value_heads=4,
    rms_norm_eps=1e-05,
    rope_theta=1000000.0,
    sliding_window=None,
    vocab_size=32000,
)

torch.manual_seed(0)
model = AutoModelForCausalLM.from_config(config)
model.eval()
device = "cpu"

batch_size = 2
seq_len = 3
past_len = 30
head_size = config.hidden_size // config.num_attention_heads
num_heads = config.num_key_value_heads
dc = transformers.cache_utils.DynamicCache()
dc.update(
    torch.rand((batch_size, num_heads, past_len, head_size)).to(device),
    torch.rand((batch_size, num_heads, past_len, head_size)).to(device),
    layer_idx=0,
)
inputs = dict(
    input_ids=torch.randint(0, 1000, (batch_size, seq_len), dtype=torch.int64).to(device),
    attention_mask=torch.randint(
        0, 1, (batch_size, seq_len + past_len), dtype=torch.int64
    ).to(device),
    past_key_values=dc,
)

############
# First wrap
############

class MistralNeMoWithKVCache(torch.nn.Module):
    def __init__(self, m):
        super().__init__()
        self.m = m

    def forward(self, input_ids, attention_mask, past_key, past_value):
        cache = transformers.cache_utils.DynamicCache()
        cache.update(past_key, past_value, layer_idx=0)
        out = self.m(
            input_ids=input_ids,
            attention_mask=attention_mask,
            past_key_values=cache,
            use_cache=True,
        )
        layer = out.past_key_values.layers[0]
        return out.logits, layer.keys, layer.values

wrapper = MistralNeMoWithKVCache(model)
wrapper.eval()

DYN = torch.export.Dim.DYNAMIC
ep = torch.export.export(
    wrapper,
    (),
    kwargs=dict(
        input_ids=inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        past_key=inputs["past_key_values"].layers[0].keys,
        past_value=inputs["past_key_values"].layers[0].values,
    ),
    dynamic_shapes=dict(
        input_ids={0: DYN, 1: DYN},
        attention_mask={0: DYN, 1: DYN},
        past_key={0: DYN, 2: DYN},
        past_value={0: DYN, 2: DYN},
    ),
)
ep.run_decompositions()

############
# Flattening
############

def _flatten_key_value_cache(cache):
    keys = [lay.keys for lay in cache.layers]
    values = [lay.values for lay in cache.layers]
    flat = list(itertools.chain.from_iterable(zip(keys, values)))
    unique = set(type(lay) for lay in cache.layers)
    assert unique == {
        transformers.cache_utils.DynamicLayer
    }, f"Not implemented for layers type {unique}"
    keys = list(
        itertools.chain.from_iterable(
            (f"key_{i}", f"value_{i}") for i in range(len(cache.layers))
        )
    )
    return flat, keys

def _flatten_with_keys_cache(cache):
    values, context = _flatten_key_value_cache(cache)
    return [(torch.utils._pytree.MappingKey(k), v) for k, v in zip(context, values)], context

def _unflatten_cache(
    values,
    context,
    output_type=None,
):
    expected = list(
        itertools.chain.from_iterable(
            (f"key_{i}", f"value_{i}") for i in range(len(values) // 2)
        )
    )
    assert expected == context, f"Does not seem to be a dynamic cache {expected} != {context}"
    res = transformers.cache_utils.DynamicCache()
    for i in range(len(values) // 2):
        res.update(values[i * 2], values[i * 2 + 1], layer_idx=i)
    assert output_type is None or isinstance(
        res, output_type
    ), f"Type mismatch between {output_type} (expected) and {type(res)}"
    return res

def registers_dynamic_cache():
    cls = transformers.cache_utils.DynamicCache
    torch.utils._pytree.register_pytree_node(
        cls,
        _flatten_key_value_cache,
        _unflatten_cache,
        serialized_type_name=f"{cls.__module__}.{cls.__name__}",
        flatten_with_keys_fn=_flatten_with_keys_cache,
    )

registers_dynamic_cache()

# checks that flattening is ok
flat_dc, spec = torch.utils._pytree.tree_flatten(inputs["past_key_values"])
assert isinstance(flat_dc, list)
assert len(flat_dc) == 2
restored = torch.utils._pytree.tree_unflatten(flat_dc, spec)
assert len(restored.layers) == 1
torch.testing.assert_close(inputs["past_key_values"].layers[0].keys, restored.layers[0].keys)
torch.testing.assert_close(
    inputs["past_key_values"].layers[0].values, restored.layers[0].values
)

ep = torch.export.export(
    model,
    (),
    kwargs=inputs,
    dynamic_shapes=dict(
        input_ids={0: DYN, 1: DYN},
        attention_mask={0: DYN, 1: DYN},
        past_key_values=[{0: DYN, 2: DYN}, {0: DYN, 2: DYN}],
    ),
)
ep.run_decompositions()  # fails here
# File "torch/export/exported_program.py", line 1530, in run_decompositions
#     return _decompose_exported_program(
#         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
# File "torch/export/exported_program.py", line 1017, in _decompose_exported_program
#     new_module_call_graph = _get_updated_module_call_graph(
#                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
# File "torch/export/exported_program.py", line 926, in _get_updated_module_call_graph
#     old_name = old_user_input_names[user_input_counter]
#             ~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^
# IndexError: list index out of range

test_mistral_nemo_torch_exoprt_export_issue()

Fix Action

Fix / Workaround

CPU: Architecture: x86_64 CPU op-mode(s): 32-bit, 64-bit Address sizes: 46 bits physical, 48 bits virtual Byte Order: Little Endian CPU(s): 20 On-line CPU(s) list: 0-19 Vendor ID: GenuineIntel Model name: 13th Gen Intel(R) Core(TM) i7-13800H CPU family: 6 Model: 186 Thread(s) per core: 2 Core(s) per socket: 10 Socket(s): 1 Stepping: 2 BogoMIPS: 5836.80 Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology tsc_reliable nonstop_tsc cpuid tsc_known_freq pni pclmulqdq vmx ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves avx_vnni vnmi umip waitpkg gfni vaes vpclmulqdq rdpid movdiri movdir64b fsrm md_clear serialize flush_l1d arch_capabilities Virtualization: VT-x Hypervisor vendor: Microsoft Virtualization type: full L1d cache: 480 KiB (10 instances) L1i cache: 320 KiB (10 instances) L2 cache: 12.5 MiB (10 instances) L3 cache: 24 MiB (1 instance) NUMA node(s): 1 NUMA node0 CPU(s): 0-19 Vulnerability Gather data sampling: Not affected Vulnerability Itlb multihit: Not affected Vulnerability L1tf: Not affected Vulnerability Mds: Not affected Vulnerability Meltdown: Not affected Vulnerability Mmio stale data: Not affected Vulnerability Reg file data sampling: Mitigation; Clear Register File Vulnerability Retbleed: Mitigation; Enhanced IBRS Vulnerability Spec rstack overflow: Not affected Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl 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 BHI_DIS_S Vulnerability Srbds: Not affected Vulnerability Tsx async abort: Not affected

Code Example

def test_mistral_nemo_torch_exoprt_export_issue():
    import itertools
    import torch
    import transformers
    from transformers import AutoModelForCausalLM, MistralConfig

    num_hidden_layers = 1
    config = MistralConfig(
        architectures=["MistralNeMoForCausalLM"],
        bos_token_id=1,
        eos_token_id=2,
        hidden_act="silu",
        hidden_size=512,
        intermediate_size=1376,
        max_position_embeddings=1024,
        model_type="mistral",
        num_attention_heads=8,
        num_hidden_layers=num_hidden_layers,
        num_key_value_heads=4,
        rms_norm_eps=1e-05,
        rope_theta=1000000.0,
        sliding_window=None,
        vocab_size=32000,
    )

    torch.manual_seed(0)
    model = AutoModelForCausalLM.from_config(config)
    model.eval()
    device = "cpu"

    batch_size = 2
    seq_len = 3
    past_len = 30
    head_size = config.hidden_size // config.num_attention_heads
    num_heads = config.num_key_value_heads
    dc = transformers.cache_utils.DynamicCache()
    dc.update(
        torch.rand((batch_size, num_heads, past_len, head_size)).to(device),
        torch.rand((batch_size, num_heads, past_len, head_size)).to(device),
        layer_idx=0,
    )
    inputs = dict(
        input_ids=torch.randint(0, 1000, (batch_size, seq_len), dtype=torch.int64).to(device),
        attention_mask=torch.randint(
            0, 1, (batch_size, seq_len + past_len), dtype=torch.int64
        ).to(device),
        past_key_values=dc,
    )

    ############
    # First wrap
    ############

    class MistralNeMoWithKVCache(torch.nn.Module):
        def __init__(self, m):
            super().__init__()
            self.m = m

        def forward(self, input_ids, attention_mask, past_key, past_value):
            cache = transformers.cache_utils.DynamicCache()
            cache.update(past_key, past_value, layer_idx=0)
            out = self.m(
                input_ids=input_ids,
                attention_mask=attention_mask,
                past_key_values=cache,
                use_cache=True,
            )
            layer = out.past_key_values.layers[0]
            return out.logits, layer.keys, layer.values

    wrapper = MistralNeMoWithKVCache(model)
    wrapper.eval()

    DYN = torch.export.Dim.DYNAMIC
    ep = torch.export.export(
        wrapper,
        (),
        kwargs=dict(
            input_ids=inputs["input_ids"],
            attention_mask=inputs["attention_mask"],
            past_key=inputs["past_key_values"].layers[0].keys,
            past_value=inputs["past_key_values"].layers[0].values,
        ),
        dynamic_shapes=dict(
            input_ids={0: DYN, 1: DYN},
            attention_mask={0: DYN, 1: DYN},
            past_key={0: DYN, 2: DYN},
            past_value={0: DYN, 2: DYN},
        ),
    )
    ep.run_decompositions()

    ############
    # Flattening
    ############

    def _flatten_key_value_cache(cache):
        keys = [lay.keys for lay in cache.layers]
        values = [lay.values for lay in cache.layers]
        flat = list(itertools.chain.from_iterable(zip(keys, values)))
        unique = set(type(lay) for lay in cache.layers)
        assert unique == {
            transformers.cache_utils.DynamicLayer
        }, f"Not implemented for layers type {unique}"
        keys = list(
            itertools.chain.from_iterable(
                (f"key_{i}", f"value_{i}") for i in range(len(cache.layers))
            )
        )
        return flat, keys

    def _flatten_with_keys_cache(cache):
        values, context = _flatten_key_value_cache(cache)
        return [(torch.utils._pytree.MappingKey(k), v) for k, v in zip(context, values)], context

    def _unflatten_cache(
        values,
        context,
        output_type=None,
    ):
        expected = list(
            itertools.chain.from_iterable(
                (f"key_{i}", f"value_{i}") for i in range(len(values) // 2)
            )
        )
        assert expected == context, f"Does not seem to be a dynamic cache {expected} != {context}"
        res = transformers.cache_utils.DynamicCache()
        for i in range(len(values) // 2):
            res.update(values[i * 2], values[i * 2 + 1], layer_idx=i)
        assert output_type is None or isinstance(
            res, output_type
        ), f"Type mismatch between {output_type} (expected) and {type(res)}"
        return res

    def registers_dynamic_cache():
        cls = transformers.cache_utils.DynamicCache
        torch.utils._pytree.register_pytree_node(
            cls,
            _flatten_key_value_cache,
            _unflatten_cache,
            serialized_type_name=f"{cls.__module__}.{cls.__name__}",
            flatten_with_keys_fn=_flatten_with_keys_cache,
        )

    registers_dynamic_cache()

    # checks that flattening is ok
    flat_dc, spec = torch.utils._pytree.tree_flatten(inputs["past_key_values"])
    assert isinstance(flat_dc, list)
    assert len(flat_dc) == 2
    restored = torch.utils._pytree.tree_unflatten(flat_dc, spec)
    assert len(restored.layers) == 1
    torch.testing.assert_close(inputs["past_key_values"].layers[0].keys, restored.layers[0].keys)
    torch.testing.assert_close(
        inputs["past_key_values"].layers[0].values, restored.layers[0].values
    )

    ep = torch.export.export(
        model,
        (),
        kwargs=inputs,
        dynamic_shapes=dict(
            input_ids={0: DYN, 1: DYN},
            attention_mask={0: DYN, 1: DYN},
            past_key_values=[{0: DYN, 2: DYN}, {0: DYN, 2: DYN}],
        ),
    )
    ep.run_decompositions()  # fails here
    # File "torch/export/exported_program.py", line 1530, in run_decompositions
    #     return _decompose_exported_program(
    #         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    # File "torch/export/exported_program.py", line 1017, in _decompose_exported_program
    #     new_module_call_graph = _get_updated_module_call_graph(
    #                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    # File "torch/export/exported_program.py", line 926, in _get_updated_module_call_graph
    #     old_name = old_user_input_names[user_input_counter]
    #             ~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^
    # IndexError: list index out of range


test_mistral_nemo_torch_exoprt_export_issue()

---

PyTorch version: 2.12.0.dev20260406+cu130
Is debug build: False
CUDA used to build PyTorch: 13.0
ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04.3 LTS (x86_64)
GCC version: (Ubuntu 13.3.0-6ubuntu2~24.04.1) 13.3.0
Clang version: Could not collect
CMake version: version 4.1.2
Libc version: glibc-2.39

Python version: 3.12.3 (main, Mar  3 2026, 12:15:18) [GCC 13.3.0] (64-bit runtime)
Python platform: Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39
Is CUDA available: True
CUDA runtime version: 13.0.88
CUDA_MODULE_LOADING set to: 
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 4060 Laptop GPU
Nvidia driver version: 581.57
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.9.14.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.14.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.14.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.14.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.14.0
/usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.14.0
/usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.14.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.14.0
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
Caching allocator config: N/A

CPU:
Architecture:                         x86_64
CPU op-mode(s):                       32-bit, 64-bit
Address sizes:                        46 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               20
On-line CPU(s) list:                  0-19
Vendor ID:                            GenuineIntel
Model name:                           13th Gen Intel(R) Core(TM) i7-13800H
CPU family:                           6
Model:                                186
Thread(s) per core:                   2
Core(s) per socket:                   10
Socket(s):                            1
Stepping:                             2
BogoMIPS:                             5836.80
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology tsc_reliable nonstop_tsc cpuid tsc_known_freq pni pclmulqdq vmx ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves avx_vnni vnmi umip waitpkg gfni vaes vpclmulqdq rdpid movdiri movdir64b fsrm md_clear serialize flush_l1d arch_capabilities
Virtualization:                       VT-x
Hypervisor vendor:                    Microsoft
Virtualization type:                  full
L1d cache:                            480 KiB (10 instances)
L1i cache:                            320 KiB (10 instances)
L2 cache:                             12.5 MiB (10 instances)
L3 cache:                             24 MiB (1 instance)
NUMA node(s):                         1
NUMA node0 CPU(s):                    0-19
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Not affected
Vulnerability Reg file data sampling: Mitigation; Clear Register File
Vulnerability Retbleed:               Mitigation; Enhanced IBRS
Vulnerability Spec rstack overflow:   Not affected
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl
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 BHI_DIS_S
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] jax2onnx==0.12.3
[pip3] mypy==1.19.1
[pip3] mypy_extensions==1.1.0
[pip3] ndonnx==0.17.2
[pip3] numpy==2.2.6
[pip3] nv-one-logger-pytorch-lightning-integration==2.3.1
[pip3] nvidia-cublas==13.1.0.3
[pip3] nvidia-cublas-cu12==12.8.4.1
[pip3] nvidia-cuda-cupti==13.0.85
[pip3] nvidia-cuda-cupti-cu12==12.8.90
[pip3] nvidia-cuda-nvrtc==13.0.88
[pip3] nvidia-cuda-nvrtc-cu12==12.8.93
[pip3] nvidia-cuda-runtime==13.0.96
[pip3] nvidia-cuda-runtime-cu12==12.8.90
[pip3] nvidia-cudnn-cu12==9.10.2.21
[pip3] nvidia-cudnn-cu13==9.20.0.48
[pip3] nvidia-cufft==12.0.0.61
[pip3] nvidia-cufft-cu12==11.3.3.83
[pip3] nvidia-curand==10.4.0.35
[pip3] nvidia-curand-cu12==10.3.9.90
[pip3] nvidia-cusolver==12.0.4.66
[pip3] nvidia-cusolver-cu12==11.7.3.90
[pip3] nvidia-cusparse==12.6.3.3
[pip3] nvidia-cusparse-cu12==12.5.8.93
[pip3] nvidia-cusparselt-cu12==0.7.1
[pip3] nvidia-cusparselt-cu13==0.8.1
[pip3] nvidia-nccl-cu12==2.27.5
[pip3] nvidia-nccl-cu13==2.29.7
[pip3] nvidia-nvjitlink==13.0.88
[pip3] nvidia-nvjitlink-cu12==12.8.93
[pip3] nvidia-nvtx==13.0.85
[pip3] nvidia-nvtx-cu12==12.8.90
[pip3] nvtx==0.2.15
[pip3] onnx==1.22.0
[pip3] onnx-ir==0.2.0
[pip3] onnx-shape-inference==0.1.7
[pip3] onnxruntime-easy==0.1.0
[pip3] onnxruntime_extensions==0.14.0
[pip3] onnxruntime-genai-cuda==0.11.4
[pip3] onnxruntime-gpu==1.24.4
[pip3] onnxslim==0.1.82
[pip3] optree==0.19.0
[pip3] pytorch-lightning==2.6.1
[pip3] torch==2.12.0.dev20260406+cu130
[pip3] torchaudio==2.11.0.dev20260402+cu130
[pip3] torchcodec==0.11.0.dev20260201+cu130
[pip3] torchmetrics==1.8.2
[pip3] torchvision==0.27.0.dev20260406+cu130
[pip3] torchx==0.7.0
[pip3] triton==3.7.0+git9c288bc5

RAW_BUFFERClick to expand / collapse

🐛 Describe the bug

This is a LLM taking a DynamicCache as input. When the model is wrapped and the flattening done in the wrapper, it works. When a flattening function is registered and no wrapper is used, it fails when calling run_decompositions on the exported program.

def test_mistral_nemo_torch_exoprt_export_issue():
    import itertools
    import torch
    import transformers
    from transformers import AutoModelForCausalLM, MistralConfig

    num_hidden_layers = 1
    config = MistralConfig(
        architectures=["MistralNeMoForCausalLM"],
        bos_token_id=1,
        eos_token_id=2,
        hidden_act="silu",
        hidden_size=512,
        intermediate_size=1376,
        max_position_embeddings=1024,
        model_type="mistral",
        num_attention_heads=8,
        num_hidden_layers=num_hidden_layers,
        num_key_value_heads=4,
        rms_norm_eps=1e-05,
        rope_theta=1000000.0,
        sliding_window=None,
        vocab_size=32000,
    )

    torch.manual_seed(0)
    model = AutoModelForCausalLM.from_config(config)
    model.eval()
    device = "cpu"

    batch_size = 2
    seq_len = 3
    past_len = 30
    head_size = config.hidden_size // config.num_attention_heads
    num_heads = config.num_key_value_heads
    dc = transformers.cache_utils.DynamicCache()
    dc.update(
        torch.rand((batch_size, num_heads, past_len, head_size)).to(device),
        torch.rand((batch_size, num_heads, past_len, head_size)).to(device),
        layer_idx=0,
    )
    inputs = dict(
        input_ids=torch.randint(0, 1000, (batch_size, seq_len), dtype=torch.int64).to(device),
        attention_mask=torch.randint(
            0, 1, (batch_size, seq_len + past_len), dtype=torch.int64
        ).to(device),
        past_key_values=dc,
    )

    ############
    # First wrap
    ############

    class MistralNeMoWithKVCache(torch.nn.Module):
        def __init__(self, m):
            super().__init__()
            self.m = m

        def forward(self, input_ids, attention_mask, past_key, past_value):
            cache = transformers.cache_utils.DynamicCache()
            cache.update(past_key, past_value, layer_idx=0)
            out = self.m(
                input_ids=input_ids,
                attention_mask=attention_mask,
                past_key_values=cache,
                use_cache=True,
            )
            layer = out.past_key_values.layers[0]
            return out.logits, layer.keys, layer.values

    wrapper = MistralNeMoWithKVCache(model)
    wrapper.eval()

    DYN = torch.export.Dim.DYNAMIC
    ep = torch.export.export(
        wrapper,
        (),
        kwargs=dict(
            input_ids=inputs["input_ids"],
            attention_mask=inputs["attention_mask"],
            past_key=inputs["past_key_values"].layers[0].keys,
            past_value=inputs["past_key_values"].layers[0].values,
        ),
        dynamic_shapes=dict(
            input_ids={0: DYN, 1: DYN},
            attention_mask={0: DYN, 1: DYN},
            past_key={0: DYN, 2: DYN},
            past_value={0: DYN, 2: DYN},
        ),
    )
    ep.run_decompositions()

    ############
    # Flattening
    ############

    def _flatten_key_value_cache(cache):
        keys = [lay.keys for lay in cache.layers]
        values = [lay.values for lay in cache.layers]
        flat = list(itertools.chain.from_iterable(zip(keys, values)))
        unique = set(type(lay) for lay in cache.layers)
        assert unique == {
            transformers.cache_utils.DynamicLayer
        }, f"Not implemented for layers type {unique}"
        keys = list(
            itertools.chain.from_iterable(
                (f"key_{i}", f"value_{i}") for i in range(len(cache.layers))
            )
        )
        return flat, keys

    def _flatten_with_keys_cache(cache):
        values, context = _flatten_key_value_cache(cache)
        return [(torch.utils._pytree.MappingKey(k), v) for k, v in zip(context, values)], context

    def _unflatten_cache(
        values,
        context,
        output_type=None,
    ):
        expected = list(
            itertools.chain.from_iterable(
                (f"key_{i}", f"value_{i}") for i in range(len(values) // 2)
            )
        )
        assert expected == context, f"Does not seem to be a dynamic cache {expected} != {context}"
        res = transformers.cache_utils.DynamicCache()
        for i in range(len(values) // 2):
            res.update(values[i * 2], values[i * 2 + 1], layer_idx=i)
        assert output_type is None or isinstance(
            res, output_type
        ), f"Type mismatch between {output_type} (expected) and {type(res)}"
        return res

    def registers_dynamic_cache():
        cls = transformers.cache_utils.DynamicCache
        torch.utils._pytree.register_pytree_node(
            cls,
            _flatten_key_value_cache,
            _unflatten_cache,
            serialized_type_name=f"{cls.__module__}.{cls.__name__}",
            flatten_with_keys_fn=_flatten_with_keys_cache,
        )

    registers_dynamic_cache()

    # checks that flattening is ok
    flat_dc, spec = torch.utils._pytree.tree_flatten(inputs["past_key_values"])
    assert isinstance(flat_dc, list)
    assert len(flat_dc) == 2
    restored = torch.utils._pytree.tree_unflatten(flat_dc, spec)
    assert len(restored.layers) == 1
    torch.testing.assert_close(inputs["past_key_values"].layers[0].keys, restored.layers[0].keys)
    torch.testing.assert_close(
        inputs["past_key_values"].layers[0].values, restored.layers[0].values
    )

    ep = torch.export.export(
        model,
        (),
        kwargs=inputs,
        dynamic_shapes=dict(
            input_ids={0: DYN, 1: DYN},
            attention_mask={0: DYN, 1: DYN},
            past_key_values=[{0: DYN, 2: DYN}, {0: DYN, 2: DYN}],
        ),
    )
    ep.run_decompositions()  # fails here
    # File "torch/export/exported_program.py", line 1530, in run_decompositions
    #     return _decompose_exported_program(
    #         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    # File "torch/export/exported_program.py", line 1017, in _decompose_exported_program
    #     new_module_call_graph = _get_updated_module_call_graph(
    #                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    # File "torch/export/exported_program.py", line 926, in _get_updated_module_call_graph
    #     old_name = old_user_input_names[user_input_counter]
    #             ~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^
    # IndexError: list index out of range


test_mistral_nemo_torch_exoprt_export_issue()

Versions

PyTorch version: 2.12.0.dev20260406+cu130
Is debug build: False
CUDA used to build PyTorch: 13.0
ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04.3 LTS (x86_64)
GCC version: (Ubuntu 13.3.0-6ubuntu2~24.04.1) 13.3.0
Clang version: Could not collect
CMake version: version 4.1.2
Libc version: glibc-2.39

Python version: 3.12.3 (main, Mar  3 2026, 12:15:18) [GCC 13.3.0] (64-bit runtime)
Python platform: Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39
Is CUDA available: True
CUDA runtime version: 13.0.88
CUDA_MODULE_LOADING set to: 
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 4060 Laptop GPU
Nvidia driver version: 581.57
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.9.14.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.14.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.14.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.14.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.14.0
/usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.14.0
/usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.14.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.14.0
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
Caching allocator config: N/A

CPU:
Architecture:                         x86_64
CPU op-mode(s):                       32-bit, 64-bit
Address sizes:                        46 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               20
On-line CPU(s) list:                  0-19
Vendor ID:                            GenuineIntel
Model name:                           13th Gen Intel(R) Core(TM) i7-13800H
CPU family:                           6
Model:                                186
Thread(s) per core:                   2
Core(s) per socket:                   10
Socket(s):                            1
Stepping:                             2
BogoMIPS:                             5836.80
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology tsc_reliable nonstop_tsc cpuid tsc_known_freq pni pclmulqdq vmx ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves avx_vnni vnmi umip waitpkg gfni vaes vpclmulqdq rdpid movdiri movdir64b fsrm md_clear serialize flush_l1d arch_capabilities
Virtualization:                       VT-x
Hypervisor vendor:                    Microsoft
Virtualization type:                  full
L1d cache:                            480 KiB (10 instances)
L1i cache:                            320 KiB (10 instances)
L2 cache:                             12.5 MiB (10 instances)
L3 cache:                             24 MiB (1 instance)
NUMA node(s):                         1
NUMA node0 CPU(s):                    0-19
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Not affected
Vulnerability Reg file data sampling: Mitigation; Clear Register File
Vulnerability Retbleed:               Mitigation; Enhanced IBRS
Vulnerability Spec rstack overflow:   Not affected
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl
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 BHI_DIS_S
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] jax2onnx==0.12.3
[pip3] mypy==1.19.1
[pip3] mypy_extensions==1.1.0
[pip3] ndonnx==0.17.2
[pip3] numpy==2.2.6
[pip3] nv-one-logger-pytorch-lightning-integration==2.3.1
[pip3] nvidia-cublas==13.1.0.3
[pip3] nvidia-cublas-cu12==12.8.4.1
[pip3] nvidia-cuda-cupti==13.0.85
[pip3] nvidia-cuda-cupti-cu12==12.8.90
[pip3] nvidia-cuda-nvrtc==13.0.88
[pip3] nvidia-cuda-nvrtc-cu12==12.8.93
[pip3] nvidia-cuda-runtime==13.0.96
[pip3] nvidia-cuda-runtime-cu12==12.8.90
[pip3] nvidia-cudnn-cu12==9.10.2.21
[pip3] nvidia-cudnn-cu13==9.20.0.48
[pip3] nvidia-cufft==12.0.0.61
[pip3] nvidia-cufft-cu12==11.3.3.83
[pip3] nvidia-curand==10.4.0.35
[pip3] nvidia-curand-cu12==10.3.9.90
[pip3] nvidia-cusolver==12.0.4.66
[pip3] nvidia-cusolver-cu12==11.7.3.90
[pip3] nvidia-cusparse==12.6.3.3
[pip3] nvidia-cusparse-cu12==12.5.8.93
[pip3] nvidia-cusparselt-cu12==0.7.1
[pip3] nvidia-cusparselt-cu13==0.8.1
[pip3] nvidia-nccl-cu12==2.27.5
[pip3] nvidia-nccl-cu13==2.29.7
[pip3] nvidia-nvjitlink==13.0.88
[pip3] nvidia-nvjitlink-cu12==12.8.93
[pip3] nvidia-nvtx==13.0.85
[pip3] nvidia-nvtx-cu12==12.8.90
[pip3] nvtx==0.2.15
[pip3] onnx==1.22.0
[pip3] onnx-ir==0.2.0
[pip3] onnx-shape-inference==0.1.7
[pip3] onnxruntime-easy==0.1.0
[pip3] onnxruntime_extensions==0.14.0
[pip3] onnxruntime-genai-cuda==0.11.4
[pip3] onnxruntime-gpu==1.24.4
[pip3] onnxslim==0.1.82
[pip3] optree==0.19.0
[pip3] pytorch-lightning==2.6.1
[pip3] torch==2.12.0.dev20260406+cu130
[pip3] torchaudio==2.11.0.dev20260402+cu130
[pip3] torchcodec==0.11.0.dev20260201+cu130
[pip3] torchmetrics==1.8.2
[pip3] torchvision==0.27.0.dev20260406+cu130
[pip3] torchx==0.7.0
[pip3] triton==3.7.0+git9c288bc5

cc @chauhang @penguinwu @avikchaudhuri @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4

extent analysis

TL;DR

The issue can be resolved by ensuring proper registration and handling of the DynamicCache in the export process, potentially by adjusting the dynamic_shapes or the registers_dynamic_cache function.

Guidance

Verify that the registers_dynamic_cache function is correctly registering the DynamicCache class with PyTorch's _pytree module.
Check the dynamic_shapes dictionary in the torch.export.export call to ensure it accurately reflects the shape of the past_key_values input.
Consider modifying the registers_dynamic_cache function to handle the DynamicCache class more robustly, potentially by adding error checking or handling edge cases.
Ensure that the torch.export.export call is using the correct version of the model and inputs to avoid any potential versioning issues.

Example

No specific code example is provided as the issue seems to be related to the interaction between PyTorch's export functionality and the custom DynamicCache class. However, the registers_dynamic_cache function could be modified to include additional error checking or handling, such as:

def registers_dynamic_cache():
    cls = transformers.cache_utils.DynamicCache
    try:
        torch.utils._pytree.register_pytree_node(
            cls,
            _flatten_key_value_cache,
            _unflatten_cache,
            serialized_type_name=f"{cls.__module__}.{cls.__name__}",
            flatten_with_keys_fn=_flatten_with_keys_cache,
        )
    except Exception as e:
        print(f"Error registering DynamicCache: {e}")

Notes

The issue seems to be related to the interaction between PyTorch's export functionality and the custom DynamicCache class. The provided code and error message suggest that there may be an issue with the way the DynamicCache is being registered or handled during the export process.

Recommendation

Apply workaround: Modify the registers_dynamic_cache function to include additional error checking or handling to ensure that the DynamicCache class is properly registered and handled during the export process. This may involve adding try-except blocks or modifying the dynamic_shapes dictionary to better reflect the shape of the past_key_values input.

Vote matrix · Quick signals

Works

Did the solution work? Tap to confirm.

Easy Fix

Was it a quick fix?

Time Saver

Did it save you time?

Blocking

Was it severely blocking?

Common Issue

Are others likely hitting this too?

Flaky / Intermittent

Is it intermittent?

Verified / Reproducible

Can you reproduce it reliably?

#api #ssr #installation #tensor shape #request timeout

Still need to ship something?

×6

Another batch ranked right after the header list — different links, same matching logic.

Data

Security

Network

Code

UI/UX

Text

System

Multimedia

Protocol

API

Engineering

pytorch - 💡(How to fix) Fix run_decomposition fails when using flattening functions [1 participants]

Recommended Tools

GitHub issue graph ai analysis

Error Message

Fix Action

Fix / Workaround

Code Example

🐛 Describe the bug

Versions

extent analysis

TL;DR

Guidance

Example

Notes

Recommendation

Still need to ship something?

TRENDING

pytorch - 💡(How to fix) Fix run_decomposition fails when using flattening functions [1 participants]

Recommended Tools

GitHub issue graph ai analysis

Error Message

Fix Action

Fix / Workaround

Code Example

🐛 Describe the bug

Versions

extent analysis

TL;DR

Guidance

Example

Notes

Recommendation

Still need to ship something?

RELATED_DISCOVERY

TRENDING