Skye Wanderman-Milne
3fe8bd027c
Adjust pmap_bechmark.py values to be more realistic. ( #2622 )
2020-04-06 16:38:34 -07:00
Skye Wanderman-Milne
c28c46e191
Add ShardedDeviceArray indexing benchmark. ( #2549 )
...
Example output:
```
---------Benchmark summary for ShardedDeviceArray_indexing---------
indices_fn mean %std relative
------------------ -------- ------- ----------
integer_indices 0.16901 8.52522 1
integer_2D_indices 18.4918 0 109.412
```
2020-03-31 15:52:41 -07:00
George Necula
fd52fbf411
Fix import in benchmarks
...
This works on my machine as 'python benchmarks/pmap_benchmark.py'. It also
follows the code in examples.
This will need a copybara rule to change the import to 'from jax.benchmarks import benchmark'
2020-03-31 11:48:08 +03:00
Skye Wanderman-Milne
24bbd2bc1d
Fix pmap_benchmark.py import ( #2524 )
2020-03-27 10:50:57 -07:00
George Necula
428377afb3
Added type annotations and removed unused imports ( #2472 )
...
* Added type annotations and removed unused imports
* Adjusted type hints for pytype
2020-03-21 13:54:30 +01:00
George Necula
cd7ab0a9e0
Changed to pmap_benchmark to make it runnable in Google ( #2448 )
2020-03-19 06:56:59 +01:00
Skye Wanderman-Milne
75077a1441
Add pmap_benchmark.py ( #2409 )
...
Example output:
```
$ TARGET_TOTAL_SECS=2 CUDA_VISIBLE_DEVICES= XLA_FLAGS=--xla_force_host_platform_device_count=500 python3 benchmarks/pmap_benchmark.py
2020-03-12 15:46:35.903121: E external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_driver.cc:313] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected
/usr/local/google/home/skyewm/jax/jax/lib/xla_bridge.py:122: UserWarning: No GPU/TPU found, falling back to CPU.
warnings.warn('No GPU/TPU found, falling back to CPU.')
---------Benchmark results for pmap_shard_args_nargs=10_nshards=4---------
mean=0.034490 std=0.002890 %std=8.378140 total=2.000426
#iters=58 #warmup=1
---------Benchmark results for pmap_shard_args_nargs=100_nshards=4---------
mean=0.091495 std=0.005935 %std=6.486871 total=2.012888
#iters=22 #warmup=1
---------Benchmark results for pmap_shard_args_nargs=101_nshards=4---------
mean=0.113549 std=0.009080 %std=7.996712 total=2.043878
#iters=18 #warmup=1
---------Benchmark results for pmap_shard_args_nargs=500_nshards=4---------
mean=0.356868 std=0.007960 %std=2.230518 total=2.141210
#iters=6 #warmup=1
---------Benchmark results for pmap_shard_args_nargs=10_nshards=2---------
mean=0.022288 std=0.002946 %std=13.219607 total=2.005951
#iters=90 #warmup=1
---------Benchmark results for pmap_shard_args_nargs=10_nshards=4---------
mean=0.035210 std=0.002024 %std=5.747389 total=2.006975
#iters=57 #warmup=1
---------Benchmark results for pmap_shard_args_nargs=10_nshards=8---------
mean=0.048641 std=0.001578 %std=3.243398 total=2.042912
#iters=42 #warmup=1
---------Benchmark results for pmap_shard_args_nargs=10_nshards=100---------
mean=0.257487 std=0.007190 %std=2.792452 total=2.059900
#iters=8 #warmup=1
---------Benchmark results for pmap_shard_args_nargs=10_nshards=500---------
mean=1.696294 std=0.005097 %std=0.300473 total=3.392588
#iters=2 #warmup=1
---------Benchmark summary for pmap_shard_args---------
nargs nshards mean %std relative
------- --------- --------- --------- ----------
10 4 0.0344901 8.37814 1
100 4 0.0914949 6.48687 2.65279
101 4 0.113549 7.99671 3.29221
500 4 0.356868 2.23052 10.347
10 2 0.0222883 13.2196 0.646224
10 4 0.0352101 5.74739 1.02088
10 8 0.0486408 3.2434 1.41028
10 100 0.257487 2.79245 7.46555
10 500 1.69629 0.300473 49.182
---------Benchmark results for pmap_shard_outputs_nouts=10_nshards=4---------
mean=0.061780 std=0.004737 %std=7.668032 total=2.038743
#iters=33 #warmup=1
---------Benchmark results for pmap_shard_outputs_nouts=100_nshards=4---------
mean=0.123264 std=0.005980 %std=4.851385 total=2.095494
#iters=17 #warmup=1
---------Benchmark results for pmap_shard_outputs_nouts=500_nshards=4---------
mean=0.471524 std=0.024051 %std=5.100792 total=2.357622
#iters=5 #warmup=1
---------Benchmark results for pmap_shard_outputs_nouts=10_nshards=2---------
mean=0.041546 std=0.004446 %std=10.700256 total=2.035745
#iters=49 #warmup=1
---------Benchmark results for pmap_shard_outputs_nouts=10_nshards=4---------
mean=0.063768 std=0.002756 %std=4.322039 total=2.040561
#iters=32 #warmup=1
---------Benchmark results for pmap_shard_outputs_nouts=10_nshards=8---------
mean=0.087285 std=0.005343 %std=6.121320 total=2.007556
#iters=23 #warmup=1
---------Benchmark results for pmap_shard_outputs_nouts=10_nshards=100---------
mean=0.623440 std=0.004038 %std=0.647725 total=2.493759
#iters=4 #warmup=1
---------Benchmark results for pmap_shard_outputs_nouts=10_nshards=500---------
mean=4.096676 std=0.000000 %std=0.000000 total=4.096676
#iters=1 #warmup=1
---------Benchmark summary for pmap_shard_outputs---------
nouts nshards mean %std relative
------- --------- --------- --------- ----------
10 4 0.0617801 7.66803 1
100 4 0.123264 4.85139 1.99521
500 4 0.471524 5.10079 7.6323
10 2 0.0415458 10.7003 0.672479
10 4 0.0637675 4.32204 1.03217
10 8 0.087285 6.12132 1.41283
10 100 0.62344 0.647725 10.0913
10 500 4.09668 0 66.3106
```
2020-03-17 14:31:25 -07:00
