# CUDA OOM when export a large model to ONNX

**URL:** <https://discuss.huggingface.co/t/cuda-oom-when-export-a-large-model-to-onnx/31877>\
**Category:** 🤗Optimum\
**Created:** [February 17, 2023, 3:08am UTC](https://discuss.huggingface.co/t/cuda-oom-when-export-a-large-model-to-onnx/31877 "2023-02-17T03:08:30Z")\
**Posts on this page:** 7\
**Page:** 1

<div class="post-metadata">

**Author:** ![in-certo](https://avatars.discourse-cdn.com/v4/letter/i/c6cbf5/32.png) [@in-certo](https://discuss.huggingface.co/u/in-certo)\
**Post date:** [February 17, 2023, 3:08am UTC](https://discuss.huggingface.co/t/cuda-oom-when-export-a-large-model-to-onnx/31877/1 "2023-02-17T03:08:30Z")

</div>

I got OOM when exporting a large model to ONNX. I wonder how Optimum handles this issue.

Here are my settings

1. The code works for a smaller model with fewer parameters, so the error is due to the model size.
2. Not able to export the model on the CPU because of fp16.
3. The pure model takes 20GB of CUDA memory, and the total GPU capacity is 80G. ( It seems 10x memory will be consumed for exporting a small model. 2GB-\>20GB)
4. Running multiple forward passes before export won’t cause any trouble.
5. A greedy search is implemented in the graph to generate 32 tokens. A lot of intermediate `past_key_values` are cached.

---

<div class="post-metadata">

**Author:** ![in-certo](https://avatars.discourse-cdn.com/v4/letter/i/c6cbf5/32.png) [@in-certo](https://discuss.huggingface.co/u/in-certo)\
**Post date:** [February 17, 2023, 8:34am UTC](https://discuss.huggingface.co/t/cuda-oom-when-export-a-large-model-to-onnx/31877/2 "2023-02-17T08:34:22Z")

</div>

Very odd. It seems that the time and memory consumed to export a jit.ScriptModule are proportional to the loop size.

If this is true, it seems impossible to export a model with a decoding method into the ONNX computation graph.

```auto
class Model2(nn.Module):
    def forward(self, x):
        for i in range(2):
            x *= x
        return x

class Model32(nn.Module):
    def forward(self, x):
        for i in range(32):
            x *= x
        return x

```

---

<div class="post-metadata">

**Author:** ![fxmarty](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/fxmarty/32/23782_2.png) [@fxmarty](https://discuss.huggingface.co/u/fxmarty)\
**Post date:** [February 17, 2023, 9:27am UTC](https://discuss.huggingface.co/t/cuda-oom-when-export-a-large-model-to-onnx/31877/3 "2023-02-17T09:27:14Z")

</div>

Thanks @in-certo , could it be linked to this issue? [`torch.jit.trace` memory usage increase although forward is constant, and gets much slower than forward with model depth increase · Issue #93943 · pytorch/pytorch · GitHub](https://github.com/pytorch/pytorch/issues/93943)

I witnessed as well the memory usage increasing with the number of loops when using torch.jit.trace with stable diffusion.

---

<div class="post-metadata">

**Author:** ![in-certo](https://avatars.discourse-cdn.com/v4/letter/i/c6cbf5/32.png) [@in-certo](https://discuss.huggingface.co/u/in-certo)\
**Post date:** [February 17, 2023, 9:33am UTC](https://discuss.huggingface.co/t/cuda-oom-when-export-a-large-model-to-onnx/31877/4 "2023-02-17T09:33:10Z")

</div>

Yes, exactly.

Here is another related issue: [ONNX model file exported from Transformer decoder is too large · Issue #4319 · onnx/onnx (github.com)](https://github.com/onnx/onnx/issues/4319)

---

<div class="post-metadata">

**Author:** ![implank](https://avatars.discourse-cdn.com/v4/letter/i/67e7ee/32.png) [@implank](https://discuss.huggingface.co/u/implank)\
**Post date:** [June 26, 2025, 1:02pm UTC](https://discuss.huggingface.co/t/cuda-oom-when-export-a-large-model-to-onnx/31877/5 "2025-06-26T13:02:39Z")

</div>

I tried to convert a PyTorch model to ONNX, but encountered an OOM error. However, using inference directly can be successful. After adding ‘with torch. reference\_made()’ before ‘torch. onnx. export’, I was able to export the model as onnx without oom

---

<div class="post-metadata">

**Author:** ![John6666](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/john6666/32/27664_2.png) [@John6666](https://discuss.huggingface.co/u/John6666)\
**Post date:** [June 26, 2025, 1:08pm UTC](https://discuss.huggingface.co/t/cuda-oom-when-export-a-large-model-to-onnx/31877/6 "2025-06-26T13:08:41Z")

</div>

For future reference…

As of today, if you want to convert a model to ONNX, it seems that using the frequently maintained Optimum library is generally recommended.

> **[Export a model to ONNX with optimum.exporters.onnx](https://huggingface.co/docs/optimum/exporters/onnx/usage_guides/export_a_model)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.

> **[Common AI Model Formats](https://huggingface.co/blog/ngxson/common-ai-model-formats)**
>
> A Blog post by Xuan-Son Nguyen on Hugging Face

---

<div class="post-metadata">

**Author:** ![surfiniaburger](https://avatars.discourse-cdn.com/v4/letter/s/f19dbf/32.png) [@surfiniaburger](https://discuss.huggingface.co/u/surfiniaburger)\
**Post date:** [July 26, 2025, 2:25pm UTC](https://discuss.huggingface.co/t/cuda-oom-when-export-a-large-model-to-onnx/31877/7 "2025-07-26T14:25:15Z")

</div>

I tried it but still got oom error `print("\nStep 2: Preparing and running the ONNX export...") try: # The `zsh: killed`error is a classic Out-Of-Memory (OOM) error from the OS. # Your insight about`torch.inference\_mode()`(you mentioned`torch.reference\_made()`,

```
# by disabling all gradient calculations during loading and exporting.
with torch.inference_mode():
    # --- Pre-load check and fix for empty index files ---
    index_path = Path(pytorch_model_path) / "model.safetensors.index.json"
    if index_path.exists() and index_path.stat().st_size == 0:
        print(f"⚠️ Found an empty index file at: {index_path}")
        print(" This can cause loading errors. Removing it to proceed.")
        os.remove(index_path)
        print(" ✅ Empty index file removed.")

    # --- MPS DEBUGGING ---
    # Forcing CPU to bypass any MPS-specific bugs.
    device = "cpu"
    print(f"Using device: {device}")

    # Load the model and config ONCE to have better control over memory.
    print("Loading model and config from disk...")
    main_config = AutoConfig.from_pretrained(pytorch_model_path, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        pytorch_model_path,
        config=main_config,
        trust_remote_code=True,
    ).to(device)
    print("Model loaded.")

    custom_onnx_configs = {
        "decoder_model": CustomGemma3NMultimodalOnnxConfig(config=main_config, task="text-generation", use_past=False),
        "decoder_with_past_model": CustomGemma3NMultimodalOnnxConfig(config=main_config, task="text-generation", use_past=True),
    }
    # Use the more direct `onnx_export_from_model` which takes a pre-loaded model object.
    onnx_export_from_model(
         model=model,
         output=Path(onnx_output_path),
         task="text-generation-with-past",
         custom_onnx_configs=custom_onnx_configs,
         fn_get_submodels=get_submodels,
         opset=14,
         do_validation=False,
         device=device,
     )
    print("\n✅ ONNX conversion process completed successfully!")
    print(f" The exported model is saved in: {Path(onnx_output_path).resolve()}")

```

except Exception:  
print(f"\n❌ An error occurred during the ONNX conversion process.“); print(”— FULL TRACEBACK —“); traceback.print\_exc(); print(”— END OF TRACEBACK —")`