(gradients, adasum=False)
return fastmath.nested_map(lambda g: g / n, gradients_psum)
Averages gradients over all the devices across different hosts.
Averages gradients over all the devices across different hosts.
Averages gradients over all the devices across different hosts .
def _average_multidevice_gradients(gradients, adasum=False):
  """Averages gradients over all the devices across different hosts."""
  n = fastmath.global_device_count() // base.N_WEIGHTS_SHARDS
  if adasum:
    # This implements a version of the Adasum algorithm from the following
    # paper: https://arxiv.org/pdf/2006.02924.pdf
    lg = max([i for i in range(20) if 2**i <= n])
    for lg_i in range(lg):
      shift = 2**lg_i
      perm = []
      for i in range(n):
        block_i = i % (2*shift)  # we do blocks of 2*shift size
        if block_i < shift:
          perm.append((i, i+shift))
        else:
          perm.append((i, i-shift))
      perm_grad = jax.lax.ppermute(gradients, perm=perm, axis_name='batch')
      gradients = fastmath.nested_map_multiarg(
          _adasum_merge, gradients, perm_grad)
  if base.N_WEIGHTS_SHARDS > 1:  # only sum gradients from matching shards
    groups = [[base.N_WEIGHTS_SHARDS * i + d for i in range(int(n))]
              for d in range(base.N_WEIGHTS_SHARDS)]
    gradients_psum = fastmath.psum(gradients, 'batch',
                                   axis_index_groups=groups)
  else:
    gradients_psum = fastmath.psum(gradients, 'batch')  # sum all gradients
  n = jnp.array(n, dtype=jnp.float32)
  return fastmath.nested_map(lambda g: g / n, gradients_psum)
def _average_multidevice_gradients ( gradients , adasum = False ) : n = fastmath . global_device_count ( ) // base . N_WEIGHTS_SHARDS if adasum : # This implements a version of the Adasum algorithm from the following # paper: https://arxiv.org/pdf/2006.02924.pdf lg = max ( [ i for i in range ( 20 ) if 2 ** i <= n ] ) for lg_i in range ( lg ) : shift = 2 ** lg_i perm = [ ] for i in range ( n ) : block_i = i % ( 2 * shift ) # we do blocks of 2*shift size if block_i < shift : perm . append ( ( i , i + shift ) ) else : perm . append ( ( i , i - shift ) ) perm_grad = jax . lax . ppermute ( gradients , perm = perm , axis_name = 'batch' ) gradients = fastmath . nested_map_multiarg ( _adasum_merge , gradients , perm_grad ) if base . N_WEIGHTS_SHARDS > 1 : # only sum gradients from matching shards groups = [ [ base . N_WEIGHTS_SHARDS * i + d for i in range ( int ( n ) ) ] for d in range ( base . N_WEIGHTS_SHARDS ) ] gradients_psum = fastmath . psum ( gradients , 'batch' , axis_index_groups = groups ) else : gradients_psum = fastmath . psum ( gradients , 'batch' ) # sum all gradients n = jnp . array ( n , dtype = jnp . float32 ) return fastmath . nested_map ( lambda g : g / n , gradients_psum )