code: import argparse
  import os
  import random
  
  import torch
  from torch import distributed as dist
  from torch.utils.data import DataLoader
  from torch.nn.parallel import DistributedDataParallel as DDP
  from torch.utils.data.distributed import DistributedSampler
  import torch.nn as nn
  import torch.nn.functional as F
  import numpy as np
  
  
  class LinModel(nn.Module):
  
      def __init__(self, in_dim, out_dim):
          super(LinModel, self).__init__()
          self.linear = nn.Linear(in_dim, out_dim)
  
      def forward(self, x):
          out = self.linear(x)
          out = F.softmax(out, dim=-1)
          return out
  
  
  def reduce_loss(tensor, rank, world_size):
      with torch.no_grad():
          dist.reduce(tensor, dst=0)
          if rank == 0:
              tensor /= world_size
  
  
  def build_fake_data(size=1000):
      x1 = [(random.uniform(0, 0.5), 0) for i in range(size // 2)]
      x2 = [(random.uniform(0.5, 1), 1) for i in range(size // 2)]
      return x1 + x2
  
  
  def evaluate(valid_loader):
      net.eval()
      with torch.no_grad():
          cnt = 0
          total = 0
          for inputs, labels in valid_loader:
              inputs, labels = inputs.unsqueeze(1).float().cuda(), labels.long().cuda()
              output = net(inputs)
              predict = torch.argmax(output, dim=1)
              cnt += (predict == labels).sum().item()
              total += len(labels)
              # print(f'right = {(predict == labels).sum()}')
          cnt = torch.Tensor([cnt]).to(inputs.device)
          total = torch.Tensor([total]).to(inputs.device)
          reduced_param = torch.cat((cnt.view(1), total.view(1)))
          cnt = reduced_param[0].item()
          total = reduced_param[1].item()
      return cnt, total
  
  
  def set_random_seed(seed):
      """Set random seed for reproducability."""
      if seed is not None and seed > 0:
          random.seed(seed)
          np.random.seed(seed)
          torch.manual_seed(seed)
          torch.cuda.manual_seed(seed)
  
  
  parser = argparse.ArgumentParser()
  parser.add_argument('--local_rank', type=int, default=-1, help="local gpu id")
  parser.add_argument('--batch_size', type=int, default=128, help="batch size")
  parser.add_argument('--lr', type=float, default=0.1, help="learn rate")
  parser.add_argument('--epochs', type=int, default=5, help="train epoch")
  parser.add_argument('--seed', type=int, default=40, help="train epoch")
  
  
  args = parser.parse_args()
  args.world_size = int(os.getenv("WORLD_SIZE", '1'))
  
  set_random_seed(args.seed)
  dist.init_process_group(backend='nccl', init_method='env://')
  torch.cuda.set_device(args.local_rank)
  global_rank = dist.get_rank()
  
  print(f'global_rank = {global_rank} local_rank = {args.local_rank} world_size = {args.world_size}')
  
  net = LinModel(1, 2)
  net.cuda()
  net = torch.nn.SyncBatchNorm.convert_sync_batchnorm(net)
  net = DDP(net, device_ids=[args.local_rank], output_device=args.local_rank)
  
  trainset = build_fake_data(size=10000)
  validset = build_fake_data(size=10000)
  
  train_sampler = DistributedSampler(trainset)
  valid_sampler = DistributedSampler(validset)
  
  train_loader = DataLoader(trainset,
                            batch_size=args.batch_size,
                            shuffle=False,
                            pin_memory=True,
                            sampler=train_sampler)
  
  valid_loader = DataLoader(validset,
                            batch_size=args.batch_size,
                            shuffle=False,
                            pin_memory=True,
                            sampler=valid_sampler)
  
  criterion = torch.nn.CrossEntropyLoss()
  opt = torch.optim.Adam(net.parameters(), lr=args.lr)
  
  net.train()
  for e in range(int(args.epochs)):
  
      for idx, (inputs, labels) in enumerate(train_loader):
          inputs = inputs.unsqueeze(1).float().cuda()
          labels = labels.long().cuda()
          output = net(inputs)
          loss = criterion(output, labels)
          opt.zero_grad()
          loss.backward()
          opt.step()
          reduce_loss(loss, global_rank, args.world_size)
          # if idx % 10 == 0 and global_rank == 0:
          #     print('Epoch: {} step: {} loss: {}'.format(e, idx, loss.item()))
      cnt, total = evaluate(valid_loader)
      if global_rank == 0:
          print(f'epoch {e} || eval accuracy: {cnt / total}')
  
  # if global_rank == 0:
  #     print(f'net weight = {net.state_dict()}')
  
apis: torch.utils.data.DataLoader torch.nn.SyncBatchNorm.convert_sync_batchnorm torch.distributed.reduce torch.nn.Linear torch.distributed.get_rank torch.utils.data.distributed.DistributedSampler torch.argmax torch.cuda.manual_seed torch.nn.functional.softmax torch.distributed.init_process_group torch.no_grad numpy.random.seed torch.manual_seed torch.nn.CrossEntropyLoss torch.nn.parallel.DistributedDataParallel torch.Tensor torch.cuda.set_device