Multiprocessing mode raise "Thread xxx creation via pthread_create failed" error

[wuciting]

I run multiple tasks in multiprocessing mode. It takes too much cpu memory, more than 50G. But my problem size is small, only 10*10.

I think it has something to do with memory preallocation? Every process tries to preallocate some memory.

Is there a environment variable to control this? I see one solution in gpu mode, but not working in my case.

I installed a cuda version jaxlib, but I set it to run in cpu mode. I think it is ok.

environment:

jax==0.4.8
jaxlib==0.4.7+cuda11.cudnn86

script:

from concurrent.futures import ProcessPoolExecutor

import numpy as np
from jax import jit, config

config.update('jax_platform_name', 'cpu')

@jit
def f(x):
    return (x ** 2).sum()

def train(n):
    return f(np.random.rand(10, 10))


if __name__ == '__main__':
    with ProcessPoolExecutor() as executor:
        for res in executor.map(train, list(range(7))):
            print(res)