Repository: dmlc/dgl
Pull request: [Bug] Fix multiple issues in distributed multi-GPU GraphSAGE example
* Previously running the GraphSAGE example with multi-GPU multi-machine will throw device mismatch error.
* Removed `pad_data` and try `torch.distributed.algorithms.join.Join` to deal with uneven training set sizes.

Problem statement:
Bug in DistDGL when the batch size is small
## 🐛 Bug

The problem appears when the used batch size is small. This results in different trainers to have different number of batches and while training some trainers will finish earlier than other. The result is that the other trainers will be stuck at synchronization.

## To Reproduce

Steps to reproduce the behavior:

1.Use small batch size.

Hints from issue discussion:
Thanks for reporting. This is also the problem for pure PyTorch distributed training, that number of batches should be the same for all the trainer, otherwise the barrier operation will be blocked. Our current solution is to manually truncate the larger num batches into the same number for all the traininer
I made a sample code at https://github.com/dmlc/dgl/pull/3542/files, to pad the train_nids. Could you check whether this fixes your problem?
@VoVAllen It looks like we had previously fixed this for NodeDataLoader/EdgeDataLoader when using multiple processes by supporting the `use_ddp` flag (https://github.com/dmlc/dgl/pull/2911), but it only applies to DGLGraph and not DistGraph?

@nv-dlasalle For distributed dataloader, we don't have support for the torch sampler, thus the solution in #2911 might not help
<strike> But we essentially need the same solution, to round up the size of the training set, so that it is evenly divisible by the number of trainers. Or perform a max-reduce across the trainers for the number of mini-batches that will be run, and have any trainers with fewer generate extra ones.</strike>

Is the plan to migrate the padding in your above change to be inside of the DataLoader in the future?
This issue has been automatically marked as stale due to lack of activity. It will be closed if no further activity occurs. Thank you
This issue is closed due to lack of activity. Feel free to reopen it if you still have questions.
> Thanks for reporting. This is also the problem for pure PyTorch distributed training, that number of batches should be the same for all the trainer, otherwise the barrier operation will be blocked. Our current solution is to manually truncate the larger num batches into the same number for all the traininer

I'm facing similar issue that the distributed training code would stuck into deadlocks "randomly" under some conditions. And when I tune the parameters like `batch_size`, `num_servers`, the behavior might change (e.g., stuck in deadlock -> run successfully). 

Is there a systematic way to solve that? Is `pad_data` the solution, just like what it does [here](https://github.com/dmlc/dgl/blob/480a4ae35c84c4497bfa901c25e6e6eca85b67eb/examples/pytorch/graphsage/dist/train_dist.py#L158) on distributed GraphSAGE? 
As @[VoVAllen](https://github.com/VoVAllen) suggested, https://pytorch.org/tutorials/advanced/generic_join.html#using-join-with-distributeddataparallel this works for this issue without `pad_data`.

Commit message: update version requirement in README

Patch to review:
diff --git a/examples/pytorch/graphsage/dist/README.md b/examples/pytorch/graphsage/dist/README.md
index 7d5b05d32e88..7d6bb4e77d2f 100644
--- a/examples/pytorch/graphsage/dist/README.md
+++ b/examples/pytorch/graphsage/dist/README.md
@@ -6,6 +6,8 @@ This is an example of training GraphSage in a distributed fashion. Before traini
 sudo pip3 install ogb
 ```
 
+**Requires PyTorch 1.10.0+ to work.**
+
 To train GraphSage, it has five steps:
 
 ### Step 0: Setup a Distributed File System
diff --git a/examples/pytorch/graphsage/dist/train_dist.py b/examples/pytorch/graphsage/dist/train_dist.py
index fcb005c5748c..163564e08afd 100644
--- a/examples/pytorch/graphsage/dist/train_dist.py
+++ b/examples/pytorch/graphsage/dist/train_dist.py
@@ -20,6 +20,8 @@
 import torch.optim as optim
 import torch.multiprocessing as mp
 from torch.utils.data import DataLoader
+from torch.distributed.algorithms.join import Join
+import socket
 
 def load_subtensor(g, seeds, input_nodes, device, load_feat=True):
     """
@@ -155,41 +157,11 @@ def evaluate(model, g, inputs, labels, val_nid, test_nid, batch_size, device):
     model.train()
     return compute_acc(pred[val_nid], labels[val_nid]), compute_acc(pred[test_nid], labels[test_nid])
 
-def pad_data(nids, device):
-    """
-    In distributed traning scenario, we need to make sure that each worker has same number of
-    batches. Otherwise the synchronization(barrier) is called diffirent times, which results in
-    the worker with more batches hangs up.
-
-    This function pads the nids to the same size for all workers, by repeating the head ids till
-    the maximum size among all workers.
-    """
-    import torch.distributed as dist
-    # NCCL backend only supports GPU tensors, thus here we need to allocate it to gpu
-    num_nodes = th.tensor(nids.numel()).to(device)
-    dist.all_reduce(num_nodes, dist.ReduceOp.MAX)
-    max_num_nodes = int(num_nodes)
-    nids_length = nids.shape[0]
-    if max_num_nodes > nids_length:
-        pad_size = max_num_nodes % nids_length
-        repeat_size = max_num_nodes // nids_length
-        new_nids = th.cat([nids for _ in range(repeat_size)] + [nids[:pad_size]], axis=0)
-        print("Pad nids from {} to {}".format(nids_length, max_num_nodes))
-    else:
-        new_nids = nids
-    assert new_nids.shape[0] == max_num_nodes
-    return new_nids
-
 
 def run(args, device, data):
     # Unpack data
     train_nid, val_nid, test_nid, in_feats, n_classes, g = data
     shuffle = True
-    if args.pad_data:
-        train_nid = pad_data(train_nid, device)
-        # Current pipeline doesn't support duplicate node id within the same batch
-        # Therefore turn off shuffling to avoid potential duplicate node id within the same batch
-        shuffle = False
     # Create sampler
     sampler = NeighborSampler(g, [int(fanout) for fanout in args.fan_out.split(',')],
                               dgl.distributed.sample_neighbors, device)
@@ -209,8 +181,7 @@ def run(args, device, data):
         if args.num_gpus == -1:
             model = th.nn.parallel.DistributedDataParallel(model)
         else:
-            dev_id = g.rank() % args.num_gpus
-            model = th.nn.parallel.DistributedDataParallel(model, device_ids=[dev_id], output_device=dev_id)
+            model = th.nn.parallel.DistributedDataParallel(model, device_ids=[device], output_device=device)
     loss_fcn = nn.CrossEntropyLoss()
     loss_fcn = loss_fcn.to(device)
     optimizer = optim.Adam(model.parameters(), lr=args.lr)
@@ -233,43 +204,46 @@ def run(args, device, data):
         # Loop over the dataloader to sample the computation dependency graph as a list of
         # blocks.
         step_time = []
-        for step, blocks in enumerate(dataloader):
-            tic_step = time.time()
-            sample_time += tic_step - start
-
-            # The nodes for input lies at the LHS side of the first block.
-            # The nodes for output lies at the RHS side of the last block.
-            batch_inputs = blocks[0].srcdata['features']
-            batch_labels = blocks[-1].dstdata['labels']
-            batch_labels = batch_labels.long()
-
-            num_seeds += len(blocks[-1].dstdata[dgl.NID])
-            num_inputs += len(blocks[0].srcdata[dgl.NID])
-            blocks = [block.to(device) for block in blocks]
-            batch_labels = batch_labels.to(device)
-            # Compute loss and prediction
-            start = time.time()
-            batch_pred = model(blocks, batch_inputs)
-            loss = loss_fcn(batch_pred, batch_labels)
-            forward_end = time.time()
-            optimizer.zero_grad()
-            loss.backward()
-            compute_end = time.time()
-            forward_time += forward_end - start
-            backward_time += compute_end - forward_end
-
-            optimizer.step()
-            update_time += time.time() - compute_end
-
-            step_t = time.time() - tic_step
-            step_time.append(step_t)
-            iter_tput.append(len(blocks[-1].dstdata[dgl.NID]) / step_t)
-            if step % args.log_every == 0:
-                acc = compute_acc(batch_pred, batch_labels)
-                gpu_mem_alloc = th.cuda.max_memory_allocated() / 1000000 if th.cuda.is_available() else 0
-                print('Part {} | Epoch {:05d} | Step {:05d} | Loss {:.4f} | Train Acc {:.4f} | Speed (samples/sec) {:.4f} | GPU {:.1f} MB | time {:.3f} s'.format(
-                    g.rank(), epoch, step, loss.item(), acc.item(), np.mean(iter_tput[3:]), gpu_mem_alloc, np.sum(step_time[-args.log_every:])))
-            start = time.time()
+
+        with Join([model]):
+            for step, blocks in enumerate(dataloader):
+                tic_step = time.time()
+                sample_time += tic_step - start
+
+                # The nodes for input lies at the LHS side of the first block.
+                # The nodes for output lies at the RHS side of the last block.
+                batch_inputs = blocks[0].srcdata['features']
+                batch_labels = blocks[-1].dstdata['labels']
+                batch_labels = batch_labels.long()
+
+                num_seeds += len(blocks[-1].dstdata[dgl.NID])
+                num_inputs += len(blocks[0].srcdata[dgl.NID])
+                blocks = [block.to(device) for block in blocks]
+                batch_labels = batch_labels.to(device)
+                # Compute loss and prediction
+                start = time.time()
+                #print(g.rank(), blocks[0].device, model.module.layers[0].fc_neigh.weight.device, dev_id)
+                batch_pred = model(blocks, batch_inputs)
+                loss = loss_fcn(batch_pred, batch_labels)
+                forward_end = time.time()
+                optimizer.zero_grad()
+                loss.backward()
+                compute_end = time.time()
+                forward_time += forward_end - start
+                backward_time += compute_end - forward_end
+
+                optimizer.step()
+                update_time += time.time() - compute_end
+
+                step_t = time.time() - tic_step
+                step_time.append(step_t)
+                iter_tput.append(len(blocks[-1].dstdata[dgl.NID]) / step_t)
+                if step % args.log_every == 0:
+                    acc = compute_acc(batch_pred, batch_labels)
+                    gpu_mem_alloc = th.cuda.max_memory_allocated() / 1000000 if th.cuda.is_available() else 0
+                    print('Part {} | Epoch {:05d} | Step {:05d} | Loss {:.4f} | Train Acc {:.4f} | Speed (samples/sec) {:.4f} | GPU {:.1f} MB | time {:.3f} s'.format(
+                        g.rank(), epoch, step, loss.item(), acc.item(), np.mean(iter_tput[3:]), gpu_mem_alloc, np.sum(step_time[-args.log_every:])))
+                start = time.time()
 
         toc = time.time()
         print('Part {}, Epoch Time(s): {:.4f}, sample+data_copy: {:.4f}, forward: {:.4f}, backward: {:.4f}, update: {:.4f}, #seeds: {}, #inputs: {}'.format(
@@ -285,11 +259,14 @@ def run(args, device, data):
                                                                                   time.time() - start))
 
 def main(args):
+    print(socket.gethostname(), 'Initializing DGL dist')
     dgl.distributed.initialize(args.ip_config)
     if not args.standalone:
+        print(socket.gethostname(), 'Initializing DGL process group')
         th.distributed.init_process_group(backend=args.backend)
+    print(socket.gethostname(), 'Initializing DistGraph')
     g = dgl.distributed.DistGraph(args.graph_name, part_config=args.part_config)
-    print('rank:', g.rank())
+    print(socket.gethostname(), 'rank:', g.rank())
 
     pb = g.get_partition_book()
     if 'trainer_id' in g.ndata:
@@ -311,7 +288,8 @@ def main(args):
     if args.num_gpus == -1:
         device = th.device('cpu')
     else:
-        device = th.device('cuda:'+str(args.local_rank))
+        dev_id = g.rank() % args.num_gpus
+        device = th.device('cuda:'+str(dev_id))
     labels = g.ndata['labels'][np.arange(g.number_of_nodes())]
     n_classes = len(th.unique(labels[th.logical_not(th.isnan(labels))]))
     print('#labels:', n_classes)


Code review comments:
diff_hunk: @@ -233,43 +204,46 @@ def run(args, device, data):
         # Loop over the dataloader to sample the computation dependency graph as a list of
         # blocks.
         step_time = []
-        for step, blocks in enumerate(dataloader):
-            tic_step = time.time()
-            sample_time += tic_step - start
-
-            # The nodes for input lies at the LHS side of the first block.
-            # The nodes for output lies at the RHS side of the last block.
-            batch_inputs = blocks[0].srcdata['features']
-            batch_labels = blocks[-1].dstdata['labels']
-            batch_labels = batch_labels.long()
-
-            num_seeds += len(blocks[-1].dstdata[dgl.NID])
-            num_inputs += len(blocks[0].srcdata[dgl.NID])
-            blocks = [block.to(device) for block in blocks]
-            batch_labels = batch_labels.to(device)
-            # Compute loss and prediction
-            start = time.time()
-            batch_pred = model(blocks, batch_inputs)
-            loss = loss_fcn(batch_pred, batch_labels)
-            forward_end = time.time()
-            optimizer.zero_grad()
-            loss.backward()
-            compute_end = time.time()
-            forward_time += forward_end - start
-            backward_time += compute_end - forward_end
-
-            optimizer.step()
-            update_time += time.time() - compute_end
-
-            step_t = time.time() - tic_step
-            step_time.append(step_t)
-            iter_tput.append(len(blocks[-1].dstdata[dgl.NID]) / step_t)
-            if step % args.log_every == 0:
-                acc = compute_acc(batch_pred, batch_labels)
-                gpu_mem_alloc = th.cuda.max_memory_allocated() / 1000000 if th.cuda.is_available() else 0
-                print('Part {} | Epoch {:05d} | Step {:05d} | Loss {:.4f} | Train Acc {:.4f} | Speed (samples/sec) {:.4f} | GPU {:.1f} MB | time {:.3f} s'.format(
-                    g.rank(), epoch, step, loss.item(), acc.item(), np.mean(iter_tput[3:]), gpu_mem_alloc, np.sum(step_time[-args.log_every:])))
-            start = time.time()
+
+        with Join([model]):
path: examples/pytorch/graphsage/dist/train_dist.py
text: @user1:
Maybe we'd better use `model.join()` as `Join` is added since torch.1.10 and `model.join()` is added prior. could you make sure of this?

@user2:
Do we need to set throw_on_early_termination=True and catch the exception?

@author:
> Do we need to set throw_on_early_termination=True and catch the exception?

`throw_on_early_termination` means throw an error if one trainer exhausted its input early.  In our case we would like to continue training.