Infinite memory accumulation over prolonged inference times

[Ethan0429]

I haven't quite figured out how yet, but it appears that there's a memory leak occurring in this library somewhere, or at least in the ONNX runtime. I've tested on multiple models so far, and the result is always the same: The program eventually accumulates enough memory to OOM, despite my efforts to batch inferences and not hold more data than necessary in memory.

My pipeline is as follows:

fn create_model(tokenizer_path: &PathBuf, model_path: &PathBuf) -> GlinerResult<GLiNER<SpanMode>> {
    let model = GLiNER::<SpanMode>::new(
        Parameters::default(),
        RuntimeParameters::default().with_threads(NUM_THREADS),
        tokenizer_path,
        model_path,
    )?;
    Ok(model)
}

Note that I have not tested this on tokenizer mode.

I'm simply reading a CSV in batches (not using the built-in CSV TextInput functionality) and creating the TextInput from it:

let input = TextInput::new(text_batch.clone(), vec!["us_organizations".to_string()])?;

I should note that my inputs are relatively large at first, and I have batches of 1,000 inputs. However I do split inputs based on their token count so that nothing more than 512 tokens gets sent to the model per input in the among the 1,000. However, I have 3,000,000 records in total. Perhaps it's not noticeable on smaller datasets, but since this pipeline runs fairly long, the memory accumulation is very clear. Each batch inferenced accumulates more and more memory.

Here is a profile of the memory during an execution (note this only processes one batch, so 1,000 records)

Am I perhaps using the library wrong, or is there a bug happening?

[fbilhaut]

Hi Ethan,

Thank you very much for this interesting report, and sorry for the late reply.

AFAIK this is not something that has been observed yet, but this would be a major issue and we definitely have to sort this out.

I've been taking care of memory lifecycle in gline-rs of course, and I don't see any problem in the snippet you provided, so I just hope the problem is not located in the ONNX runtime. But we'll see...

I will do some tests on my side, could you please indicate what is your particular configuration (namely the OS, Rust version, and ort/GPU settings if not running on CPU) ?

Also If you could send the flamegraph in SVG format it would be useful.

Thanks !

[elcuervo]

Interesting to see this.
I'm building a Ruby gem that is using this library as an extension and I noticed the same issue through it.
An unbound memory growth. Originally I though it was related to my extension code but it might be related to this scenario.

Also using SpanMode running on CPU

[Ethan0429]

@fbilhaut Apologies for the late reply, I had only just now seen your response. Unfortunately, I no longer have access to either the code nor the SVG flamegraph for this particular issue. It was part of a project I was working on (solo) at a company I no longer work for. Though I would still be interested in seeing this problem solved! My suspicion is that it lies within the ONNX runtime, but that's not based off much. Unfortunately I have neither the compute nor the resources at home to reproduce the issue effectively, but perhaps @elcuervo could chime in with their specific snippet/diagnostics.

I will however add some information:

• Rhel 8
• CPU, 40 threads
• SpanMode
• Rust version 1.89.0

[fbilhaut]

Hi @elcuervo, @Ethan0429,

Thanks a lot for adding those extra details.

I’ve just added a few basic tools for memory profiling, to trace both the memory allocated directly by the Rust code itself, and the overall process memory usage (which, of course, also includes what’s consumed by the ONNX runtime).

I ran them with the CPU benchmark in span mode, with the following findings:

• The memory used on the Rust side doesn’t change by a single byte from one inference to the next (according to the measurements of the cap crate, which instruments the allocator). This tends to confirm that there’s no memory leak in the Rust code itself (which would have been surprising anyway, given the guarantees the language provides at compile time in this area).
• The overall memory usage of the process isn’t perfectly stable on my setup, but I haven’t been able to observe any clear upward or downward trend. (However, it seems to be less stable than in token mode.)

I therefore suspect that what you’re seeing is indeed related to an issue on the ONNX runtime side, possibly tied to a specific version, system, and/or model…

@elcuervo would it be possible for you to run a test with the updated branch by running the CPU benchmark? You would just need to pull the main branch, update the path to the model in examples/benchmark_cpu.rs and enable the memprof feature:

cargo run --features memprof --release --example benchmark-cpu

You will probably want to increase the values of the REPEAT and/or MAX_SAMPLES constants to put more load on the benchmark.

I’ll continue running more extensive tests on my side, but your feedback would be greatly appreciated (along with your system specs and the model you’re using).

Thanks again!

[fbilhaut]

Some more testing today, on another setup, and I still cannot reproduce any increase in memory usage. Here is what I observed:

• The memory allocated in Rust still perfectly steady (byte-for-byte).
• Global usage of the process varies a little before stabilizing after something like 100,000 samples. After that the global usage remains steady as well (byte for byte).
• In the end I had more than 250K samples processed without any further evolution.

Here are the details:

• gline-rs: 1.0.0, in span mode, on CPU, with load-dynamic
• Model: onnx-community/gliner_small-v2.1
• OS: Windows 11
• ONNX Runtime: 1.20.1 (win-x86-gpu)
• Rust: 1.88.0
• Batch size: 500 sentences (from NuNER)
• Threads: 24

250K sample is not that much, and they are always the same for each batch, but nothing is supposed to be cached and I believe that a memory leak happening in this configuration would have been clearly noticeable.

@elcuervo if you are able to reproduce your problem with the same bench on your dataset, it would be very interesting to know the details of your setup (especially OS, ONNX runtime version, and model, and batch sizes).