opened
EiffL
username_0
This issue is to track the development of a benchmark script to test the scaling of a distributed FFT in Mesh TensorFlow under different environments.

I have a prototype script for Cori in https://github.com/username_1/flowpm/tree/mesh/scripts

But there is a lot of room for improvement, in particular:

  - [ ] Measuring the communication and compute times
  - [ ] Support TPU enviroments as well as SLURM cluster
  - [ ] Automatically run a scaling experiment and report the results

Help is most welcome :-)
Building FFT benchmark script
issue
created
EiffL
username_0
Adding a PR to support exporting timeline to profile the execution of FFTs, look at this beautiful graph of flows between the 16 GPUs of a 1D mesh on 2 cori-gpu nodes: 
![image](https://user-images.githubusercontent.com/861591/68538104-0d981c80-0324-11ea-9118-9d429d7e5fc0.png)
comment
created
EiffL
username_0
We now have an FFT benchmark that runs on both TPUs and Slurm cluster. And the first results are super interesting!

The scripts running these benchmarks are here: https://github.com/username_1/flowpm/tree/mesh/scripts
All it does is sample a batch of random 3D cubes i.e. tensors of size 128x512x512x512, execute a forward/backward FFT transform, and check accuracy. This runs on a 1D mesh, with the array splitted along the first spatial dimension.

Here are 2 traces captured on Cori on a one vs two nodes for the same volume size:
  - 1 Node: [timeline_1_node.log](https://github.com/username_1/flowpm/files/3829345/timeline_1_node.log)
  - 2 Nodes: [timeline_2_nodes.log](https://github.com/username_1/flowpm/files/3829346/timeline_2_nodes.log)
To see these traces, simply go to Chrome [tracing plugin](about:tracing) and load these files.

Here, the computation on 2 nodes is about 10x slower than on a single node. It is certainly expected that cross-nodes communications are going to be the bottleneck, but we are using completely default TensorFlow with gRPC at the moment. There is a lot of room for improvement there.

Also ran a trace on a TPU slice with the same code essentially. There we get very interesting numbers and insight into our very naive first implementation of 3D FFTs
  - Dominated by all-to-all communications, 56% of the time, and 10% of the time doing memory operations, during array transpose:
![image](https://user-images.githubusercontent.com/861591/68555326-12bba100-03e2-11ea-84a1-da63eb5cae9f.png)
FFTs account for 18% of the time, and even there they only use a tiny amount of the compute.
  - It looks like we can improve on the amount of communication by a factor 2, by combining the transpose and reshape steps we are currently using to transpose and split the memory array between two FFTs:
![image](https://user-images.githubusercontent.com/861591/68555639-84e0b580-03e3-11ea-9ffa-d225073b3fe8.png)
comment
created
EiffL
username_0
@username_1 can you just report on your latest experiments. and then we can tick the last box and close this issue :-)
comment
created
modichirag
username_1
Timings for a single FFT+iFFT for different grids and configurations of mesh on 1, 2, 4, 8 nodes of GPUs on Cori

![image](https://user-images.githubusercontent.com/13356766/74394232-0a984680-4dc1-11ea-9005-6d8926adbad2.png)

And this is the comparison with timings for the same operation in the differentiable python code ([vmad](https://github.com/rainwoodman/vmad)) run on Cori Haswell ( 1 and 2 nodes respectively as required by the number of processes)

![image](https://user-images.githubusercontent.com/13356766/74394162-e9cff100-4dc0-11ea-8f2f-941905a5fff9.png)
comment
closed
EiffL
username_0
issue