text: import os 
  import torch
  import torch.distributed as dist
  from torch.nn import SyncBatchNorm
  import copy
  
  def link_dist(func):
      def wrapper(*args, **kwargs):
          dist_init()
          func(*args, **kwargs)
          dist_finalize()
      return wrapper
  
  def dist_init(method='slurm', device_id=0):
      if method == 'slurm':
          proc_id = int(os.environ['SLURM_PROCID'])
          # ntasks = int(os.environ['SLURM_NTASKS'])
          # node_list = os.environ['SLURM_NODELIST']
          num_gpus = torch.cuda.device_count()
          torch.cuda.set_device(proc_id % num_gpus)
      elif method == 'single_node':
          torch.cuda.set_device(device_id)
  
      initialize()
      world_size = dist.get_world_size()
      rank = dist.get_rank()
      return rank, world_size
  
  def dist_finalize():
      pass
  
  def initialize(backend='nccl'):
      port = "36799"
      proc_id = int(os.environ['SLURM_PROCID'])
      ntasks = int(os.environ['SLURM_NTASKS'])
      node_list = os.environ['SLURM_NODELIST']
      if '[' in node_list:
          beg = node_list.find('[')
          pos1 = node_list.find('-', beg)
          if pos1 < 0:
              pos1 = 1000
          pos2 = node_list.find(',', beg)
          if pos2 < 0:
              pos2 = 1000
          node_list = node_list[:min(pos1, pos2)].replace('[', '')
      addr = node_list[8:].replace('-', '.')
      os.environ['MASTER_PORT'] = port
      os.environ['MASTER_ADDR'] = addr
      os.environ['WORLD_SIZE'] = str(ntasks)
      os.environ['RANK'] = str(proc_id)
      if backend == 'nccl':
          dist.init_process_group(backend='nccl')
      else:
          dist.init_process_group(backend='gloo', rank=proc_id, world_size=ntasks)
      rank = dist.get_rank()
      device = rank % torch.cuda.device_count()
      torch.cuda.set_device(device)
  
  class DistModule(torch.nn.Module):
      def __init__(self, module, sync=False):
          super(DistModule, self).__init__()
          self.module = module
          self.broadcast_params()
  
          self.sync = sync
          if not sync:
              self._grad_accs = []
              self._register_hooks()
  
      def forward(self, *inputs, **kwargs):
          return self.module(*inputs, **kwargs)
  
      def _register_hooks(self):
          for i, (name, p) in enumerate(self.named_parameters()):
              if p.requires_grad:
                  p_tmp = p.expand_as(p)
                  grad_acc = p_tmp.grad_fn.next_functions[0][0]
                  grad_acc.register_hook(self._make_hook(name, p, i))
                  self._grad_accs.append(grad_acc)
  
      def _make_hook(self, name, p, i):
          def hook(*ignore):
              # dist.allreduce_async(name, p.grad.data)
              dist.all_reduce(name, p.grad.data)
          return hook
  
      def sync_gradients(self):
          """ average gradients """
          if self.sync and dist.get_world_size() > 1:
              for name, param in self.module.named_parameters():
                  if param.requires_grad and param.grad is not None:
                      dist.all_reduce(param.grad.data)
          else:
              dist.synchronize()
  
      def broadcast_params(self):
          """ broadcast model parameters """
          for name, param in self.module.state_dict().items():
              dist.broadcast(param, 0)
  
  
  