Distributed.all_reduce bandwidth expectations

username: 98f9a893
cooked: <p>I want to benchmark how quickly PyTorch with the Gloo backend is able to all-reduce all-gather a model synchronously. To do so, I’ve written the following script [2] working with the latest Gloo backend / PyTorch. I start it on N machines, and then they together all-reduce it fine. However, the bandwidth that I see, irrespective of N, is 0.5 * link_bandwidth. The N machines are all connected to a 100 Mbps per-port switch. This is expected with a large N, as the documentation does state that it uses a ring all-reduce/all-gather to perform the distributed.all_reduce behind that scenes.</p>
<p>However, for i.e. N=2 I would expect it to perform the all-reduce at 1.0 * link_bandwidth, as one node only needs to send to one other node its full model. My experiments show however 0.5 * link_bandwidth [3]. I would expect the bandwidth of the all-reduce for any N to be [1]:</p>
<p>Amount of data to send in all-reduce: (N - 1) / N<br>
Amount of data to send in all-gather: (N - 1) / N<br>
Total data for each node to send: 2 * (N - 1) / N<br>
I.e., at N=2 we only need to send 2 * 1/2 = 1x model, and at N-&gt;inf we have to send 2x model. Translating into a all-reduce/all-gather bandwidth of 1.0 * link_bandwidth (N=2) and 0.5 * link_bandwidth (N-&gt;inf).</p>
<p>I am not sure where my calculation is wrong, or where I am misunderstanding the all-reduce ring method.</p>
<p>[1] <a href="https://github.com/zhangruiskyline/DeepLearning/blob/master/doc/system.md#allreduce-in-practice" rel="nofollow noopener">https://github.com/zhangruiskyline/DeepLearning/blob/master/doc/system.md#allreduce-in-practice</a></p>
<p>[2] allreduce.py –  execute for N=2 on each of the two machines: <code>python allreduce.py 0 2 192.168.0.1 10000</code> and <code>python allreduce.py 1 2 192.168.0.1 10000</code></p>
<pre><code class="lang-auto">#!/usr/bin/env python

import os
import sys
import torch
import torch.distributed as dist
from torch.multiprocessing import Process
import time

# Values are 4 bytes each, so 2 * 1000 * 1000 * 32 * 4 = 256 MB = 2048 Mbit
MODEL_SIZE_VALUES = 2 * 1000 * 1000 * 32
BIT_PER_VALUE = 4 * 8
BITS_PER_MBIT = 1000 * 1000


def current_time_in_ms():
    return int(round(time.time() * 1000))


def run(rank, size):
    group = dist.new_group(list(range(size)))
    tensor = torch.ones(MODEL_SIZE_VALUES, dtype=torch.float32)
    print("Performing allreduce...")
    print("   &gt; Data to send: %d Mbit" % ((MODEL_SIZE_VALUES * BIT_PER_VALUE) / float(BITS_PER_MBIT)))
    start = current_time_in_ms()
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=group)
    elapsed_ms = current_time_in_ms() - start
    print("   &gt; Finished.")
    print("   &gt; Time: %.2f s" % (elapsed_ms / 1000.0))
    print("   &gt; Speed: %.2f Mbit/s" % ((MODEL_SIZE_VALUES * BIT_PER_VALUE / BITS_PER_MBIT) / float(elapsed_ms / 1000.0)))
    print('   &gt; Result: Rank ', rank, ' has data ', str(tensor), '.\n')


def init_process(my_rank, size, master_address, master_port, fn, backend='gloo'):
    # Initialize the distributed environment
    os.environ['MASTER_ADDR'] = master_address
    os.environ['MASTER_PORT'] = master_port

    # Initialize process group
    print("Initializing process group...")
    dist.init_process_group(backend, rank=my_rank, world_size=size)
    print("   &gt; Initialized.")
    print("")
    fn(my_rank, size)


def main(my_rank, size, master_address, master_port):
    p = Process(target=init_process, args=(my_rank, size, master_address, master_port, run))
    p.start()
    p.join()


if __name__ == "__main__":
    args = sys.argv[1:]
    if len(args) != 4:
        print("Usage: python allreduce.py &lt;my rank&gt; &lt;size&gt; &lt;master address&gt; &lt;master port&gt;")
        exit(1)
    else:
        main(int(args[0]), int(args[1]), str(args[2]), str(args[3]))
</code></pre>
<p>[3] Output of machine 1 and 2 with N=2:</p>
<p>Machine 1:</p>
<pre><code class="lang-auto">Initializing process group...
   &gt; Initialized.

Performing allreduce...
   &gt; Data to send: 2048 Mbit
   &gt; Finished.
   &gt; Time: 44.79 s
   &gt; Speed: 45.72 Mbit/s
   &gt; Result: Rank  0  has data  tensor([2., 2., 2.,  ..., 2., 2., 2.]) .
</code></pre>
<p>Machine 2:</p>
<pre><code class="lang-auto">Initializing process group...
   &gt; Initialized.

Performing allreduce...
   &gt; Data to send: 2048 Mbit
   &gt; Finished.
   &gt; Time: 44.79 s
   &gt; Speed: 45.72 Mbit/s
   &gt; Result: Rank  1  has data  tensor([2., 2., 2.,  ..., 2., 2., 2.]) .
</code></pre>

username: pietern
cooked: <p>Thanks for posting the question. Your analysis is correct, it should be equal to link bandwidth for N=2 (provided the input dimensions are large enough for the runtime to be dominated by the bandwidth). The implementation in Gloo ensures there is always 1 chunk in flight and 1 chunk being reduced, so it is possible there is something wrong there for the N=2 case.</p>
<p>I’ll investigate and post back here.</p>
<p>edit: I added <a href="https://github.com/facebookincubator/gloo/issues/169" rel="nofollow noopener">https://github.com/facebookincubator/gloo/issues/169</a> to not lose track of it.</p>

username: pietern
cooked: <p>I got around to running a test with the Gloo benchmark tool and confirmed the issue.</p>
<p>For a larger explanation of the issue and the fix, see <a href="https://github.com/facebookincubator/gloo/pull/192" rel="nofollow noopener">https://github.com/facebookincubator/gloo/pull/192</a>.</p>
<p>Once this is merged and bumped in PyTorch, I expect you’ll be able to run the same test and find the bandwidth to be very close to link speed.</p>