Jaxley parallelization/speedup questions

[ewesteinde]

Hi Jaxley group, I've been working with Jaxley to try to optimize single neuron models and have noticed that the simulation time increases substantially with any increase in simulation length, compartment #, as well with batch size with vmapped simulations. Unexpectedly vmapped simulations on CPU have also been faster than the same over a GPU. The faster CPU runtime I assume is because my batch size is below what would be expected to show an improvement with a GPU, but I assumed the runtime increases I've been seeing on the CPU was due to resource limitations. However when I look at my overall CPU usage during training it consistently sits below 10%. I have a 20 CPU core system and during a simulation a small subset show variable %usage between 1-30 % and a single core sits much higher, near 100%. So I am hoping that if I can improve the parallelization across cores that it would speed up the runtime, especially when trying to simulate multiple stimuli simultaneously, but independently.

Are there strategies compatible with the Jaxley framework (especially with the Cell object class) beyond jit and vmap that you would recommend I try to improve simulation speed and/or multi-core recruitment? Some things I looked at suggested that using the Jax.pmap function (instead of vmap) would help, but it doesn't seem to be compatible with the Jaxley Cell class.

Thanks!

Some numbers:

On CPU running a vmapped simulation of batch size 4 of a cell with 165 compartments and 16 trainable parameters takes about 0.05s for a 50 ms simulation at 0.1 ms sampling rate. To do the same but with a simulation time of 1.4 seconds it takes about 6.5s. Running these simulations on CPU is roughly 3x faster than running them on a NVIDIA GPU (with vmapping).

Simulation code:

def simulate(cell, params, current_stim, dt):
    current_stim = cell.branch(0).loc(0.0).data_stimulate(current = current_stim)
    return jx.integrate(cell, params=params, delta_t=dt, data_stimuli=current_stim, voltage_solver="jaxley.thomas")
    
 batched_sim = vmap(jit(simulate, static_argnums=(0,3)), in_axes=(None, None, 0, None))

[michaeldeistler]

Hi Elena,

thanks for reaching out!

The increase in runtime

Jaxley (and vmap) does indeed not automatically parallelize across CPU cores. The standard way of doing this in JAX is to use pmap (see below).

pmap not working

For me, pmap works if I keep the cell in global state, instead of marking it as static:

# Set `xla_force_host_platform_device_count` to the number of CPU cores you want to use. 
# If you are running this in a jupyter notebook, add the following line `os.environ`.
# If you are in a Python script, do: XLA_FLAGS='--xla_force_host_platform_device_count =8'  python3 your_file.py
import os
os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'

cell = ...
dt = ...

def simulate(params, current_stim):
    current_stim = cell.branch(0).loc(0.0).data_stimulate(current = current_stim)
    return jx.integrate(cell, params=params, delta_t=dt, data_stimuli=current_stim, voltage_solver="jaxley.thomas")

batched_sim = pmap(jit(simulate), in_axes=(None, 0))

Does this work for you?

The runtimes you report

I am quite confused about the superlinear increase in computation time when the simulation takes longer. One idea to fix the supralinear behavior: You currently do vmap(jit(...)). Could you try jit(vmap(...)) in your simulations?

I tried reproducing your supralinear runtime increase on my MacBook Pro (CPU), but could not: For a cell with 198 compartments, batch size 4, dt=0.1, I get:

• 50 ms simulation time: 0.015 seconds
• 1400 ms simulation time: 0.42 seconds

Could you try running the following code? (SWC file is attached)

import time
import jax.numpy as jnp
import jaxley as jx
from jax import vmap, jit
from jaxley.channels import HH

cell = jx.read_swc("morph_allen_485574832.txt", ncomp=2)
cell.insert(HH())
cell.soma.branch(0).comp(0).record()
cell.make_trainable("HH_gNa")
parameters = cell.get_parameters()

def simulate(cell, params, current_stim, dt):
    current_stim = cell.branch(0).loc(0.0).data_stimulate(current = current_stim)
    return jx.integrate(cell, params=params, delta_t=dt, data_stimuli=current_stim, voltage_solver="jaxley.thomas")

batched_sim = vmap(jit(simulate, static_argnums=(0,3)), in_axes=(None, None, 0, None))

# Simulation time step and batch size n.
dt = 0.1
n = 4

t_max = 50.0

current_stims = jx.datapoint_to_step_currents(0.0, t_max, jnp.linspace(0, 0.2, n), dt, t_max)
start_time = time.time()
v = batched_sim(cell, parameters, current_stims, dt).block_until_ready()
print("Compile: ", time.time() - start_time)  # for me: 12.950

current_stims = jx.datapoint_to_step_currents(0.0, t_max, jnp.linspace(0.01, 0.21, n), dt, t_max)
start_time = time.time()
v = batched_sim(cell, parameters, current_stims, dt).block_until_ready()
print("Run: ", time.time() - start_time)  # for me: 0.0149

t_max = 1400.0

current_stims = jx.datapoint_to_step_currents(0.0, t_max, jnp.linspace(0, 0.2, n), dt, t_max)
start_time = time.time()
v = batched_sim(cell, parameters, current_stims, dt).block_until_ready()
print("Compile: ", time.time() - start_time)  # for me: 13.586

current_stims = jx.datapoint_to_step_currents(0.0, t_max, jnp.linspace(0.01, 0.21, n), dt, t_max)
start_time = time.time()
v = batched_sim(cell, parameters, current_stims, dt).block_until_ready()
print("Run: ", time.time() - start_time)  # for me: 0.4257

morph_allen_485574832.txt

[ewesteinde]

Hi Michael,
Thanks for the fast reply! Setting cell as a global variable did get pmap working, which resulted in a ~5x speedup of my batched simulation function compared to using jit and vmap! And sorry for the confusion on the runtime numbers, they were output by my training code which I had forgotten called the simulate function twice per step, separately for a train and val batch, which was behind the larger than expected increase. Running the code you shared I get similar run times, longer compile but that's likely due to a system difference (linux desktop).

CPU run 1:
Compile: 23.9597008228302
Run: 0.016837596893310547
Compile: 23.939589500427246
Run: 0.4185445308685303

CPU run 2 (switched the order of vmap and jit when defining batched_sim)
Compile: 23.675302505493164
Run: 0.016275644302368164
Compile: 24.123381853103638
Run: 0.4244229793548584

GPU run 1:
Compile: 5.23017144203186
Run: 0.24605250358581543
Compile: 11.775845289230347
Run: 6.868589162826538

[michaeldeistler]

Setting cell as a global variable did get pmap working, which resulted in a ~5x speedup of my batched simulation function compared to using jit and vmap!

Great! I have never used pmap myself, so please let us know if anything does not work as expected.

Happy coding!
Michael

[michaeldeistler]

Regarding the compilation time: Yes, I think this is likely a system difference. We are working on a new solver for the voltages which will (hopefully) largely get down compilation times, see here.