title: Gradient Accumulation with JAX
body: <p>I made a simple script to try to do gradient accumulation with JAX. The idea is to have large batch size (e.g. 64) that are split in small chunks (e.g. 4) that fit in the GPU's memory. For each chunck, the resulting gradient, stored in a pytree, is added to the current batch gradient. The update is done only when all chunks of the large batch are computed. In this particular example, we simply try to fit random 512-dimensional vectors to random booleans with a linear layer. Here is the script:</p>
  <pre><code>import jax
  import jax.numpy as jnp
  from jax import jit, random
  from jax.experimental import optimizers
  from functools import partial
  from jax.nn.initializers import normal, zeros
  from typing import Callable
  from dataclasses import dataclass
  
  @dataclass
  class Jax_model:
      init_fun: Callable
      apply_fun: Callable
  
  
  def Dense(input_size: int, output_size: int, init_kernel=normal(), init_bias=zeros):
  
      def init_fun(key):
          key, sub_key1, sub_key2 = jax.random.split(key, 3)
          params = {
              'I': init_kernel(sub_key1, (input_size, output_size) ),
              'I_b': init_bias(sub_key2, (1,output_size) ),
          }
          return params
  
      def apply_fun(params, inputs):
          I, I_b, = params['I'], params['I_b']
          logits = inputs @ I + I_b
          return logits
  
      return Jax_model(init_fun, apply_fun)
  
  
  def divide_pytree(pytree, div):
      for pt in jax.tree_util.tree_leaves(pytree):
          pt = pt / div
      return pytree
  
  
  def add_pytrees(pytree1, pytree2):
      for pt1, pt2 in zip( jax.tree_util.tree_leaves(pytree1), jax.tree_util.tree_leaves(pytree2) ):
          pt1 = pt1 + pt2
      return pytree1
  
  
  rng_key = random.PRNGKey(42)
  batch_size = 64
  accumulation_size = 4
  model_dim = 512
  n_iter = 50
  
  model = Dense(model_dim, 1)
  rng_key, sub_key = random.split(rng_key)
  init_params = model.init_fun(sub_key)
  opt_init, opt_update, get_params = optimizers.adam(0.001)
  opt_state = opt_init(init_params)
  
  @jit
  def update(i, current_opt_state, current_batch):
      N = current_batch[0].shape[0]
      K = accumulation_size
      num_gradients = N//K
      accumulation_batch = (current_batch[ib][0:K] for ib in range(len(current_batch)))
      value, grads = jax.value_and_grad(loss_func)(get_params(current_opt_state), accumulation_batch)
      value = value / num_gradients
      grads = divide_pytree(grads, num_gradients)
      for k in range(K,N,K):
          accumulation_batch = (current_batch[ib][k:k+K] for ib in range(len(current_batch)))
          new_value, new_grads = jax.value_and_grad(loss_func)(get_params(current_opt_state), accumulation_batch)
          value = value + (new_value / num_gradients)
          grads = add_pytrees(grads, divide_pytree(new_grads, num_gradients))
      return opt_update(i, grads, current_opt_state), value
  
  def loss_func(current_params, current_batch):
      inputs, labels = current_batch
      predictions = model.apply_fun(current_params, inputs)
      loss = jnp.square(labels-predictions).sum()
      return loss
  
  for i in range(n_iter):
      rng_key, sub_key1, sub_key2 = random.split(rng_key, 3)
      inputs = jax.random.uniform(sub_key1, (batch_size, model_dim))
      labels = jax.random.uniform(sub_key2, (batch_size, 1)) &gt; 0.5
      batch = inputs, labels
      opt_state, batch_loss = update(i, opt_state, batch)
      print(i, batch_loss)
  </code></pre>
  <p>I have doubts about the <code>divide_pytree</code> and <code>add_pytrees</code>. Does it actually modify the current batch gradient or am I missing something ? Moreover, do you see any speed issue with this code ? In particular, should I use the <code>jax.lax.fori_loop</code> in place of the traditional python for loop ?</p>
  <p>Related links:</p>
  <ul>
  <li><a href="https://github.com/google/jax/issues/1488" rel="nofollow noreferrer">https://github.com/google/jax/issues/1488</a></li>
  <li><a href="https://github.com/google-research/long-range-arena/issues/4" rel="nofollow noreferrer">https://github.com/google-research/long-range-arena/issues/4</a></li>
  </ul>