Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Memory leak? #14

Open
Ebanflo42 opened this issue Mar 26, 2024 · 0 comments
Open

Memory leak? #14

Ebanflo42 opened this issue Mar 26, 2024 · 0 comments

Comments

@Ebanflo42
Copy link

I have the following minimal reproducible example, which clearly eats up CPU memory when I run it (same issue with GPU client):

use anyhow::Result;
extern crate xla;
use xla::ArrayElement;

fn main() -> Result<()> {
    let client = xla::PjRtClient::cpu()?;
    let builder = xla::XlaBuilder::new("test");
    let x = builder.parameter(0, f32::TY, &[1000000], "x")?;
    let exec = x.build()?.compile(&client)?;
    let mut x_val = xla::Literal::vec1(&[0f32; 1000000]);
    loop {
        let xla_buffers = exec.execute(&[&x_val])?;
        let literal_out = xla_buffers[0][0].to_literal_sync()?;
        x_val = literal_out.clone();
    }
}

In principle, x_val is like the weights of a network being updated at each execution. Of course, it is very much not optimal to be transferring them to and from the GPU at every iteration, but I would still expect this to not leak memory, since the custom drop functions are being called every time the buffers/literals are going out of scope. But maybe there is something drastically incorrect about how I am using the library?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant