guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-05-16 04:41:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Update for jaxDecomp pure JAX

Browse source
This commit is contained in:
Wassim KABALAN 2024-08-07 23:52:13 +02:00
parent 831291c1f9
commit 2ea05a1cd6
9 changed files with 214 additions and 532 deletions
Download patch file Download diff file Expand all files Collapse all files

18
benchmarks/run_all_jobs.sh
View file

@ -1,19 +1,21 @@
#!/bin/bash
# Run all slurms jobs
nodes_v100=(1 2 4 8 16)
nodes_a100=(1 2 4 8 16)
nodes_v100=(1 2 4 8 16 32)
nodes_a100=(1 2 4 8 16 32)
for n in ${nodes_v100[@]}; do
sbatch --nodes=$n --job-name=v100_$n-JAXPM particle_mesh_v100.slurm
sbatch --account=tkc@v100 --nodes=$n --gres=gpu:4 --tasks-per-node=4 -C v100-32g --job-name=JAXPM-$n-N-v100 particle_mesh.slurm
done
for n in ${nodes_a100[@]}; do
sbatch --nodes=$n --job-name=a100_$n-JAXPM particle_mesh_a100.slurm
sbatch --account=tkc@a100 --nodes=$n --gres=gpu:4 --tasks-per-node=4 -C a100 --job-name=JAXPM-$n-N-a100 particle_mesh.slurm
done
# single GPUs
sbatch --job-name=JAXPM-1GPU-V100 --nodes=1 --gres=gpu:1 --tasks-per-node=1 particle_mesh_v100.slurm
sbatch --job-name=JAXPM-1GPU-A100 --nodes=1 --gres=gpu:1 --tasks-per-node=1 particle_mesh_a100.slurm
sbatch --job-name=PMWD-v100 pmwd_v100.slurm
sbatch --job-name=PMWD-a100 pmwd_a100.slurm
sbatch --account=tkc@a100 --nodes=1 --gres=gpu:1 --tasks-per-node=1 -C a100 --job-name=JAXPM-1GPU-V100 particle_mesh.slurm
sbatch --account=tkc@v100 --nodes=1 --gres=gpu:1 --tasks-per-node=1 -C v100-32g --job-name=JAXPM-1GPU-A100 particle_mesh.slurm
sbatch --account=tkc@a100 --nodes=1 --gres=gpu:1 --tasks-per-node=1 -C a100 --job-name=PMWD-1GPU-v100 pmwd_pm.slurm
sbatch --account=tkc@v100 --nodes=1 --gres=gpu:1 --tasks-per-node=1 -C v100-32g --job-name=PMWD-1GPU-a100 pmwd_pm.slurm