Closed RobertTLange closed 3 years ago
Also test different ways to limit the threads/CPU usage of JAX. There seems to be an ongoing discussion how to this the best way. Something along these lines:
os.environ['XLA_FLAGS'] = f"--xla_force_host_platform_device_count={num_cpus}"
os.environ["XLA_FLAGS"] += os.pathsep + ("--xla_cpu_multi_thread_eigen=false "
"intra_op_parallelism_threads=1")
Addressed in 51d4233. Probably have to do some changes once I get to setup the GCP VM instances.
Add a helper function in order to setup environment variables. This includes memory preallocation for GPUs, CPU cores for threading/mp, number of visible devices and TPU XLA bridge setup. So something along these lines:
Potentially also add visible device option here or separately. The jax setup function should in principle be executed whenever in
log_config
themodel_type
is set to JAX, but we can leave it also optional.