text: import torch
  import torch.distributed as dist
  from torch.nn.modules import Module
  from torch.autograd import Variable
  
  def _flatten_dense_tensors(tensors):
      """Flatten dense tensors into a contiguous 1D buffer. Assume tensors are of
      same dense type.
      Since inputs are dense, the resulting tensor will be a concatenated 1D
      buffer. Element-wise operation on this buffer will be equivalent to
      operating individually.
      Arguments:
          tensors (Iterable[Tensor]): dense tensors to flatten.
      Returns:
          A contiguous 1D buffer containing input tensors.
      """
      if len(tensors) == 1:
          return tensors[0].contiguous().view(-1)
      flat = torch.cat([t.contiguous().view(-1) for t in tensors], dim=0)
      return flat
  
  def _unflatten_dense_tensors(flat, tensors):
      """View a flat buffer using the sizes of tensors. Assume that tensors are of
      same dense type, and that flat is given by _flatten_dense_tensors.
      Arguments:
          flat (Tensor): flattened dense tensors to unflatten.
          tensors (Iterable[Tensor]): dense tensors whose sizes will be used to
            unflatten flat.
      Returns:
          Unflattened dense tensors with sizes same as tensors and values from
          flat.
      """
      outputs = []
      offset = 0
      for tensor in tensors:
          numel = tensor.numel()
          outputs.append(flat.narrow(0, offset, numel).view_as(tensor))
          offset += numel
      return tuple(outputs)
  
  def apply_gradient_allreduce(module):
      if not hasattr(dist, '_backend'):
          module.warn_on_half = True
      else:
          module.warn_on_half = True if dist._backend == dist.dist_backend.GLOO else False
  
      for p in module.state_dict().values():
          if not torch.is_tensor(p):
              continue
          dist.broadcast(p, 0)
  
      def allreduce_params():
          if (module.needs_reduction):
              module.needs_reduction = False
              buckets = {}
              for param in module.parameters():
                  if param.requires_grad and param.grad is not None:
                      tp = param.data.dtype
                      if tp not in buckets:
                          buckets[tp] = []
                      buckets[tp].append(param)
              if module.warn_on_half:
                  if torch.cuda.HalfTensor in buckets:
                      print("WARNING: gloo dist backend for half parameters may be extremely slow." +
                            " It is recommended to use the NCCL backend in this case. This currently requires" +
                            "PyTorch built from top of tree master.")
                      module.warn_on_half = False
  
              for tp in buckets:
                  bucket = buckets[tp]
                  grads = [param.grad.data for param in bucket]
                  coalesced = _flatten_dense_tensors(grads)
                  dist.all_reduce(coalesced)
                  coalesced /= dist.get_world_size()
                  for buf, synced in zip(grads, _unflatten_dense_tensors(coalesced, grads)):
                      buf.copy_(synced)
  
      for param in list(module.parameters()):
          def allreduce_hook(*unused):
              Variable._execution_engine.queue_callback(allreduce_params)
  
          if param.requires_grad:
              param.register_hook(allreduce_hook)
  
      def set_needs_reduction(self, input, output):
          self.needs_reduction = True
  
      module.register_forward_hook(set_needs_reduction)
      return module
  
  def reduce_tensor(tensor, n_gpus):
      rt = tensor.clone()
      dist.all_reduce(rt, op=dist.ReduceOp.SUM)
      rt /= n_gpus
      return rt