Unexpected memory leak while sharding.

[aayushdesai]

Hello,

I am using sharding to parallelise the computation of my function over 8 CPU cores.

Here is the code:

import jax.numpy as jnp
import numpy as np
import os
from jax import jit,random
import jax
from jax.sharding import PositionalSharding
from jax.experimental import mesh_utils

def lomb_scargle_jax(t, y, frequency):

    t = t.reshape(-1,1)
    y = y.reshape(-1,1)
    omega = 2 * jnp.pi * frequency
    y_mean = jnp.mean(y)
    y_var = jnp.var(y)
    
    # Compute tau
    tau = jnp.arctan2(jnp.sum(jnp.sin(2 * omega * t)), jnp.sum(jnp.cos(2 * omega * t))) / (2 * omega)
    
    # Compute periodogram components
    cos_term = jnp.cos(omega * (t - tau))
    sin_term = jnp.sin(omega * (t - tau))
    P = (jnp.sum((y - y_mean) * cos_term,axis=0)**2 / jnp.sum(cos_term**2,axis=0) +
         jnp.sum((y - y_mean) * sin_term,axis=0)**2 / jnp.sum(sin_term**2,axis=0)) / (2 * y_var)
    
    return P



device_count = 8
bootstrap_len = 10
spacing = 400000
samples = 2500

os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={device_count}'
print(jax.local_device_count(),device_count)
sharding = PositionalSharding(mesh_utils.create_device_mesh((8,)))
jit_ls = jit(lomb_scargle_jax)

mag = 18 + random.normal(random.PRNGKey(0),shape=(samples,bootstrap_len))
mag = mag.reshape(-1)
time_LS = np.repeat(np.linspace(0,500,samples),bootstrap_len).reshape(-1)
omega = np.repeat(np.ones(spacing),bootstrap_len).reshape(-1)

omega_shard = jax.device_put(jnp.array(omega), sharding.reshape(8))
time_shard = jax.device_put(jnp.array(time_LS), sharding.replicate())
mag_shard = jax.device_put(jnp.array(mag), sharding.replicate())

print(mag_shard.shape,time_shard.shape)
# print(mag_shard.shape)
jax.debug.visualize_array_sharding(mag_shard)
jax.debug.visualize_array_sharding(time_shard)
jax.debug.visualize_array_sharding(omega_shard)
powers = jit_ls(time_shard,mag_shard,omega_shard)
print("Finished Power Computation",powers.shape)
powers = powers * 5
print("Multiplied by 5",powers.shape)

A basic run-through of the code:

1. I initialise 3 arrays (meant to be an observable, the corresponding time series, and a frequency grid).
2. I run a period-finding algorithm called Lomb Scargle.

The problem:
I choose to shard the frequency(omega) array equally across all my 8 cores and replicate the observable and time arrays across my 8 cores. In total I have 2500 points in my time series (time and observable) and I have 400,000 points in my frequency grid. This is one calculation which I want to evaluate 10 (bootstrap_len) times in parallel.

At the end I should have the resulting array power of size 4,000,000.

When I use the memory profiler on my MacBook, the memory demand peaks at 20 GBs before the code crashes. From a rough calculation I don't expect the calculation to take more than a couple of 100 MBs. What am I missing?

Finally, the print statement at line 57 evaluates successfully, but as soon as I try to compute (or even access the first element of the power array) further with the power array, the system fails.

What am I missing?

System information:
Macbook Pro 2020 M1 chip
8 GB ram

Packages:
Jax version: 0.4.25
numpy version: 1.26.4

Thank you very much.