Sure. The code is attached below. The GPU is NVIDIA V100, and GPUDirect is used.
To run the code,

a dummy fixed-size sample dataset has to be generated. Sample size is 3 * 224 * 224 * 4 Byte. The attached script can be used to generate this dataset.


#! /bin/bash
base="base.file"
dataset_base="your_dir_path"

truncate -s 602112 $base

for class in {0..9}
do
  dir="$dataset_base/${class}"
  /bin/rm -rf $dir
  mkdir -p $dir
  echo $dir created
  for img_id in {0..1300}
  do
    fpath="${dataset_base}/${class}/${img_id}.fake"
    cp $base $fpath
  done
done




You have to set master addr and port as env variable, and change the root path in vgg11.py to the created dir path


we have two files (one vgg, one data_loader).


My test case: 1 GPU per node, 16 nodes, 128 samples per GPU


python vgg11.py [batch_size] [rank] [rank_sizes]


The print out  at the end of the output is the allreduce time spent for applying allreduce directly on each parameter.
vgg11.py


import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist

import os
import sys
import time

import data_loader


cfg = {
    'VGG11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'VGG13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'VGG16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'VGG19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}

class VGG(nn.Module):
    def __init__(self, vgg_name):
        super(VGG, self).__init__()
        self.features = self.large_make_layers(cfg[vgg_name])
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 4096),
            nn.ReLU(),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(),
            nn.Dropout(),
            nn.Linear(4096, 1000),
        )

    def forward(self, x):
        out = self.features(x)
        out = out.view(out.size(0), -1)
        out = self.classifier(out)
        return out
    def large_make_layers(self, cfg, batch_norm=False):
      layers = []
      in_channels = 3
      for v in cfg:
          if v == 'M':
              layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
          else:
              conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
              if batch_norm:
                  layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU()]
              else:
                  layers += [conv2d, nn.ReLU()]
              in_channels = v
      return nn.Sequential(*layers)


def sync_gradients(model, batch_idx, timer):
  """ Gradient averaging. """
  global record_event_cnt

  for param in model.parameters():
    print(param.grad.data.shape)
    print("record_event_cnt: %d, batch_idx: %d" % (record_event_cnt, batch_idx))

    if batch_idx > 0:
      put_timer(record_event_cnt, 1, timer)
    dist.all_reduce(param.grad.data, op=dist.ReduceOp.SUM)

    if batch_idx > 0:
      put_timer(record_event_cnt, 0, timer)
      record_event_cnt += 1
def cal_single_time(count, timer):
  tot_time = 0.0
  global para_cnt
  for i in range(count):
    print("i: %d" % (i))
    time = timer[i].elapsed_time(timer[i + para_cnt])
    print("%d time: %lf" % (i, time))

def put_timer(i, start, timer):
  global para_cnt
  if i >= 0:
    if start == 1:
      timer[i].record()
      print("put start for %d " % (i))
    elif start == 0:
      #print("put timer for iteration: " + str(i))
      timer[para_cnt + i].record()
      print("put end for %d" % (para_cnt + i))



N = int(sys.argv[1])
rank = int(sys.argv[2])
world_size = int(sys.argv[3])

record_event_cnt = 0

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # PyTorch device
dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)

net = VGG('VGG11')
net = net.cuda()
#net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[0])

para_cnt = 22
sync_timer=[]
for i in range(para_cnt * 2):
  sync_timer.append(torch.cuda.Event(enable_timing=True))


start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)

optimizer = optim.SGD(net.parameters(), lr=0.1, momentum=0.9)
criterion = nn.CrossEntropyLoss()

#inputs = torch.ones([N, 3, 224, 224], device=device)
#labels = torch.empty(N, dtype=torch.long, device=device).random_(1000)
root_path='your_dir_path'
res_size=224
trainset = data_loader.DatasetFolder(root=root_path, loader=data_loader.raw_data_loader, \
    img_size=res_size, extensions=data_loader.IMG_EXTENSIONS, transform=None)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=N, shuffle=True, num_workers=1, pin_memory=True)
torch.cuda.synchronize()


for batch_idx, (inputs, labels) in enumerate(trainloader):
  inputs, labels = inputs.to(device), labels.to(device)

  if batch_idx == 1:
    torch.cuda.synchronize()
    start = time.time()
    start_event.record()

  out = net(inputs)
  loss = criterion(out, labels)

  loss.backward()
  sync_gradients(net, batch_idx, sync_timer)
  print("================")

  optimizer.step()
  optimizer.zero_grad()

  if batch_idx == 2:
    break

  record_event_cnt = 0

end_event.record()
torch.cuda.synchronize()
end = time.time()
print(end-start)
print("iter:%d, %d: %lf, cuda time: %lf"% (batch_idx, N, (end - start), start_event.elapsed_time(end_event)))
print("end record_event_cnt: %d"% (record_event_cnt))
cal_single_time(record_event_cnt, sync_timer)


data_loader.py (some codes are borrowed from original torchvision data loader)
from torchvision import datasets, transforms
import torch
import torchvision
import os
import os.path
import time
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # PyTorch device


def raw_data_loader(path, size, d):
  file_content = torch.from_file(path, dtype=torch.float, size=size)
  #file_content = file_content.to(torch.float)
  file_content.resize_((3, d, d))
  return file_content


IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp', '.fake')

def has_file_allowed_extension(filename, extensions):
    """Checks if a file is an allowed extension.
    Args:
        filename (string): path to a file
        extensions (tuple of strings): extensions to consider (lowercase)

    Returns:
        bool: True if the filename ends with one of given extensions
    """
    return filename.lower().endswith(extensions)
def make_dataset(directory, class_to_idx, extensions=None, is_valid_file=None):
    instances = []
    directory = os.path.expanduser(directory)
    both_none = extensions is None and is_valid_file is None
    both_something = extensions is not None and is_valid_file is not None
    if both_none or both_something:
        raise ValueError("Both extensions and is_valid_file cannot be None or not None at the same time")
    if extensions is not None:
        def is_valid_file(x):
            return has_file_allowed_extension(x, extensions)
    for target_class in sorted(class_to_idx.keys()):
        class_index = class_to_idx[target_class]
        target_dir = os.path.join(directory, target_class)
        if not os.path.isdir(target_dir):
            continue
        for root, _, fnames in sorted(os.walk(target_dir, followlinks=True)):
            for fname in sorted(fnames):
                path = os.path.join(root, fname)
                if is_valid_file(path):
                    item = path, class_index
                    instances.append(item)
    return instances

class DatasetFolder(datasets.VisionDataset):
  def __init__(self, root, loader, img_size, extensions=None, transform=None,
                 target_transform=None, is_valid_file=None):
        super(DatasetFolder, self).__init__(root, transform=transform,
                                            target_transform=target_transform)
        classes, class_to_idx = self._find_classes(self.root)
        samples = make_dataset(self.root, class_to_idx, extensions, is_valid_file)
        if len(samples) == 0:
            msg = "Found 0 files in subfolders of: {}\n".format(self.root)
            if extensions is not None:
                msg += "Supported extensions are: {}".format(",".join(extensions))
            raise RuntimeError(msg)

        self.loader = loader
        self.extensions = extensions

        self.img_size = img_size * img_size * 3
        self.img_res = img_size

        self.classes = classes
        self.class_to_idx = class_to_idx
        self.samples = samples
        self.targets = [s[1] for s in samples]
  def _find_classes(self, dir):
        """
        Finds the class folders in a dataset.

        Args:
            dir (string): Root directory path.

        Returns:
            tuple: (classes, class_to_idx) where classes are relative to (dir), and class_to_idx is a dictionary.

        Ensures:
            No class is a subdirectory of another.
        """
        classes = [d.name for d in os.scandir(dir) if d.is_dir()]
        classes.sort()
        class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
        return classes, class_to_idx

  def __getitem__(self, index):
        """
        Args:
            index (int): Index

        Returns:
            tuple: (sample, target) where target is class_index of the target class.
        """
        path, target = self.samples[index]
        sample = self.loader(path, self.img_size, self.img_res)
        if self.transform is not None:
            sample = self.transform(sample)
        if self.target_transform is not None:
            target = self.target_transform(target)

        return sample, target


  def __len__(self):
        return len(self.samples)