code: def multi_node_mean_grad(self, model, zero_fill=False):
      self._init_comms()
      stream = chainer.cuda.Stream.null
  
      params = _memory_utility.extract_params_set_grad(model, zero_fill)
      itemsize = 4
      n_elems_total = _memory_utility.count_grad_elements(
          params, zero_fill
      )
      n_elems_per_node = int(math.ceil(n_elems_total / self.inter_size))
      n_elems_buffer = n_elems_per_node * self.inter_size
      n_bytes_per_node = n_elems_per_node * itemsize
      n_bytes_buffer = n_bytes_per_node * self.inter_size
  
      self.gpu_buffer_a.assign(n_bytes_buffer)
      self.gpu_buffer_b.assign(n_bytes_buffer)
  
      allreduce_grad_dtype = np.float32
  
      self._pack_params_to_buffer(
          params,
          "grad",
          buffer=self.gpu_buffer_a,
          allreduce_grad_dtype=allreduce_grad_dtype,
          zero_fill=zero_fill,
      )
  
      if chainer.is_debug():
          stream.synchronize()
          array_a = self.gpu_buffer_a.array(n_elems_total)
          array_b = self.gpu_buffer_b.array(n_elems_total)
          self._check_ready_to_allreduce(array_a, array_b)
  
      # Intra-node reduce
      self.intra_nccl_comm.reduce(
          self.gpu_buffer_a.ptr(),
          self.gpu_buffer_b.ptr(),
          n_elems_total,
          nccl.NCCL_FLOAT,
          nccl.NCCL_SUM,
          0,
          stream.ptr,
      )
  
      # Inter-node allreduce
      if self.intra_rank == 0:
          self.cpu_buffer_a.assign(n_bytes_buffer)
          self.cpu_buffer_b.assign(n_bytes_buffer)
  
          arr_b = self.gpu_buffer_b.array(n_elems_buffer)
          arr_b.data.copy_to_host(
              self.cpu_buffer_b.ptr(), n_bytes_buffer
          )
  
          self.inter_mpi_comm.Alltoall(
              [
                  self.cpu_buffer_b.buffer(n_bytes_buffer),
                  mpi4py.MPI.FLOAT,
              ],
              [
                  self.cpu_buffer_a.buffer(n_bytes_buffer),
                  mpi4py.MPI.FLOAT,
              ],
          )
  
          # Reduction in GPU
          arr_a = self.gpu_buffer_a.array(n_elems_buffer)
          arr_a.data.copy_from_host(
              self.cpu_buffer_a.ptr(), n_bytes_buffer
          )
          arr_a = arr_a.reshape(self.inter_size, n_elems_per_node)
          arr_a = arr_a.sum(axis=0)
          arr_a *= 1.0 / self.size
          arr_a.data.copy_to_host(
              self.cpu_buffer_a.ptr(), n_bytes_per_node
          )
  
          self.inter_mpi_comm.Allgather(
              [
                  self.cpu_buffer_a.buffer(n_bytes_per_node),
                  mpi4py.MPI.FLOAT,
              ],
              [
                  self.cpu_buffer_b.buffer(n_bytes_buffer),
                  mpi4py.MPI.FLOAT,
              ],
          )
  
          arr_b.data.copy_from_host(
              self.cpu_buffer_b.ptr(), n_bytes_buffer
          )
  
      # Intra-node bcast
      self.intra_nccl_comm.bcast(
          self.gpu_buffer_b.ptr(),
          n_elems_total,
          nccl.NCCL_FLOAT,
          0,
          stream.ptr,
      )
  
      if chainer.is_debug():
          stream.synchronize()
          self._ensure_all_finite(
              self.gpu_buffer_b.array(n_elems_total)
          )
  
      self._unpack_params_from_buffer(
          params,
          "grad",
          self.gpu_buffer_b,
          allreduce_grad_dtype,
          zero_fill,
      )
  
file_path: chainer/chainermn/communicators/non_cuda_aware_communicator.py
imports: import warnings
  import chainer.cuda
  import chainerx
  import math
  import mpi4py.MPI
  import numpy
  from chainermn.communicators import _communication_utility
  from chainermn.communicators import _memory_utility
  from chainermn.communicators import mpi_communicator_base
  from chainermn import nccl
label_desc: Use image node_id 5 for calling the NonCudaAwareCommunicator obj's underlying member method code with example usage: obj.multi_node_mean_grad(model, zero_fill) without return types