Merge pull request #127 from borisdayma/pjit-t5x

feat(train): pjit optimization and distributed shampoo support

src/dalle_mini/data.py

@@ -152,24 +152,15 @@ class Dataset:
                     ),
                 )
 
-    def dataloader(
-        self, split, per_device_batch_size, gradient_accumulation_steps=None, epoch=None
-    ):
-        num_devices = jax.local_device_count()
-
+    def dataloader(self, split, batch_size, epoch=None):
         def _dataloader_datasets_non_streaming(
             dataset: Dataset,
-            per_device_batch_size: int,
-            gradient_accumulation_steps: int,
             rng: jax.random.PRNGKey = None,
         ):
             """
             Returns batches of size `batch_size` from truncated `dataset`, sharded over all local devices.
             Shuffle batches if rng is set.
             """
-            batch_size = (
-                per_device_batch_size * num_devices * gradient_accumulation_steps
-            )
             steps_per_epoch = len(dataset) // batch_size
 
             if rng is not None:
@@ -185,18 +176,10 @@ class Dataset:
             for idx in batch_idx:
                 batch = dataset[idx]
                 batch = {k: jnp.array(v) for k, v in batch.items()}
-                if gradient_accumulation_steps is not None:
-                    batch = jax.tree_map(
-                        lambda x: x.reshape((-1, per_device_batch_size) + x.shape[1:]),
-                        batch,
-                    )
                 yield batch
 
         def _dataloader_datasets_streaming(
             dataset: Dataset,
-            split: str,
-            per_device_batch_size: int,
-            gradient_accumulation_steps: int,
             epoch: int,
         ):
             keys = ["input_ids", "attention_mask", "labels", "decoder_input_ids"]
@@ -208,28 +191,15 @@ class Dataset:
                 # For validation data we put the entire set on each host as we could lose
                 # too many samples on pods
                 if epoch is not None:
-                    # reshuffle training data at each epoch (not applicable with validation set)
+                    assert split == "train"
+                    # reshuffle training data at each epoch
                     dataset.set_epoch(epoch)
                     epoch += 1
                 for item in dataset:
                     for k, v in item.items():
                         batch[k].append(v)
-                        # batch = 5, devices = 8, accumulation = 2 / batch_size = 5 x 8
-                        # (40, 3, 3) -> shard 8 x (5, 3, 3)
-                        # (16, 5, 3, 3) -> shard 8 x (2, 5, 3, 3)
-                    if len(batch[keys[0]]) == per_device_batch_size * num_devices * (
-                        gradient_accumulation_steps
-                        if gradient_accumulation_steps is not None
-                        else 1
-                    ):
+                    if len(batch[keys[0]]) == batch_size:
                         batch = {k: jnp.array(v) for k, v in batch.items()}
-                        if gradient_accumulation_steps is not None:
-                            batch = jax.tree_map(
-                                lambda x: x.reshape(
-                                    (-1, per_device_batch_size) + x.shape[1:]
-                                ),
-                                batch,
-                            )
                         yield batch
                         batch = {k: [] for k in keys}
                 first_loop = False
@@ -242,15 +212,11 @@ class Dataset:
             raise ValueError(f'split must be "train" or "eval", got {split}')
 
         if self.streaming:
-            return _dataloader_datasets_streaming(
-                ds, split, per_device_batch_size, gradient_accumulation_steps, epoch
-            )
+            return _dataloader_datasets_streaming(ds, epoch)
         else:
             if split == "train":
                 self.rng_dataset, input_rng = jax.random.split(self.rng_dataset)
-            return _dataloader_datasets_non_streaming(
-                ds, per_device_batch_size, gradient_accumulation_steps, input_rng
-            )
+            return _dataloader_datasets_non_streaming(ds, input_rng)
 
     @property
     def length(self):

src/dalle_mini/model/modeling.py

@@ -312,7 +312,7 @@ class FlaxBartPreTrainedModel(FlaxBartPreTrainedModel):
         seed: int = 0,
         dtype: jnp.dtype = jnp.float32,
         abstract_init: bool = False,
-        load_on_cpu: bool = True,
+        load_on_cpu: bool = False,
         **kwargs,
     ):
         module = self.module_class(config=config, dtype=dtype, **kwargs)

tools/train/distributed_shampoo.py

@@ -1,7 +1,5 @@
-"""File copied from https://github.com/google-research/google-research/edit/master/scalable_shampoo/optax/distributed_shampoo.py"""
-
 # coding=utf-8
-# Copyright 2021 The Google Research Authors.
+# Copyright 2022 The Google Research Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -147,6 +145,12 @@ class QuantizedValue:
         return val
 
 
+@struct.dataclass
+class TrainingMetrics:
+    inverse_pth_root_errors: chex.Array  # Error for inverse-pth roots.
+    # TODO(rohananil): Add more important metrics to track during training.
+
+
 # Per parameter optimizer state used in data-parallel training.
 class ParameterStats(NamedTuple):
     """State associated to each parameter of the model being trained."""
@@ -156,6 +160,7 @@ class ParameterStats(NamedTuple):
     preconditioners: List[Any]  # Preconditioners (QuantizedValue, chex.Array)
     diagonal_momentum: QuantizedValue  # Momentum for the diagonal preconditioner
     momentum: QuantizedValue  # Momentum for the shampoo preconditioner
+    training_metrics: TrainingMetrics  # Metrics (optional for training).
 
 
 # For training extremely large model; We keep a global state with a concatenated
@@ -166,6 +171,7 @@ class ParameterStats(NamedTuple):
 class GlobalShardedParameterStats:
     statistics: chex.Array  # Statistics
     preconditioners: chex.Array  # Preconditioners
+    exponents: chex.Array  # exponents
 
 
 # These are per-parameter local states; All statistics here mirror the parameter
@@ -177,12 +183,34 @@ class LocalShardedParameterStats:
     diagonal_statistics: QuantizedValue  # Accumulator for diagonal preconditioner
     diagonal_momentum: QuantizedValue  # Momentum for the diagonal preconditioner
     momentum: QuantizedValue  # Momentum for the shampoo preconditioner
+    training_metrics: TrainingMetrics  # Metrics (optional for training).
     index_start: np.int32 = struct.field(
         pytree_node=False
     )  # Index into global statistics array
     sizes: Any = struct.field(pytree_node=False)  # Sizes of the statistics.
 
 
+def init_training_metrics(num_statistics):
+    if num_statistics:
+        return TrainingMetrics(jnp.zeros([num_statistics], jnp.float32))
+    else:
+        return TrainingMetrics([])
+
+
+def init_training_metrics_shapes(num_statistics):
+    if num_statistics:
+        return TrainingMetrics([[num_statistics], jnp.float32])
+    else:
+        return TrainingMetrics([None, jnp.float32])
+
+
+def init_training_metrics_pspec(num_statistics):
+    if num_statistics:
+        return TrainingMetrics(pjit.PartitionSpec())
+    else:
+        return TrainingMetrics(None)
+
+
 class ShardedShampooStats(NamedTuple):
     """Shampoo state in sharded mode."""
 
@@ -195,6 +223,12 @@ class ShampooState(NamedTuple):
     stats: Any
 
 
+class InitFnState(NamedTuple):
+    init_fn: Any
+    pspec_fn: Any
+    shape_and_dtype_fn: Any
+
+
 class GraftingType(enum.IntEnum):
     SGD = 1
     ADAGRAD = 2
@@ -292,6 +326,8 @@ def matrix_inverse_pth_root(
       matrix^(-1/p)
     """
 
+    assert matrix.shape[0] == matrix.shape[1]
+
     # We use float32 for the matrix inverse pth root.
     # Switch to f64 if you have hardware that supports it.
     matrix_size = matrix.shape[0]
@@ -615,6 +651,7 @@ def _convert_to_parameter_stats(global_stats, local_stat):
         new_preconditioners,
         local_stat.diagonal_momentum,
         local_stat.momentum,
+        local_stat.training_metrics,
     )
 
 
@@ -624,11 +661,40 @@ def _convert_from_parameter_stats(parameter_stats, local_stats):
         parameter_stats.diagonal_statistics,
         parameter_stats.diagonal_momentum,
         parameter_stats.momentum,
+        parameter_stats.training_metrics,
         local_stats.index_start,
         local_stats.sizes,
     )
 
 
+def _add_error_into_local_stats(local_stats, errors, inverse_failure_threshold):
+    """Adds errors back into local statistics."""
+    new_local_stats = []
+    for local_stat in local_stats:
+        index_start = int(local_stat.index_start)
+        index_end = int(len(local_stat.sizes)) + index_start
+        per_stat_error = errors[index_start:index_end]
+        if local_stat.sizes:
+            per_stat_error = jnp.where(
+                jnp.logical_and(
+                    per_stat_error > 0.0, per_stat_error != inverse_failure_threshold
+                ),
+                per_stat_error,
+                local_stat.training_metrics.inverse_pth_root_errors,
+            )
+        new_local_stats.append(
+            LocalShardedParameterStats(
+                local_stat.diagonal_statistics,
+                local_stat.diagonal_momentum,
+                local_stat.momentum,
+                TrainingMetrics(per_stat_error),
+                local_stat.index_start,
+                local_stat.sizes,
+            )
+        )
+    return new_local_stats
+
+
 def batch(x, num_devices):
     """Batch `x` so that so that leading axis is num_devices."""
     n = len(x)
@@ -670,7 +736,8 @@ def distributed_shampoo(
     batch_axis_name=None,
     ### Only set following 3 params in pjit/spmd mode.
     ### WARNING: Experimental
-    mesh_axis_names=None,
+    statistics_partition_spec=None,
+    preconditioner_partition_spec=None,
     num_devices_for_pjit=None,
     shard_optimizer_states=False,
     ###
@@ -730,7 +797,8 @@ def distributed_shampoo(
       exponent_override: Override the exponent used in matrix inverse.
       batch_axis_name: labeled axis over pmap for data-parallel training the
         optimizer used for.
-      mesh_axis_names: Axis names for the mesh (used in pjit).
+      statistics_partition_spec: PartitionSpec to be used in sharded mode.
+      preconditioner_partition_spec: PartitionSpec to be used in sharded mode.
       num_devices_for_pjit: Number of devices to parallelize over when using pjit.
       shard_optimizer_states: Shard optimizer states to save memory in model
         parallel training.
@@ -830,6 +898,11 @@ def distributed_shampoo(
         )
 
     def sharded_init_fn(params):
+        """Returns optimizer state (for PJIT mode).
+
+        Args:
+          params: the parameters that should be updated.
+        """
         params_flat, treedef = jax.tree_flatten(params)
         # Find max size to pad to.
         max_size = 0
@@ -845,6 +918,7 @@ def distributed_shampoo(
         padded_statistics = []
         padded_preconditioners = []
         local_stats_flat = []
+        exponents = []
         for param in params_flat:
             preconditioner = Preconditioner(
                 param, block_size, best_effort_shape_interpretation
@@ -862,6 +936,12 @@ def distributed_shampoo(
                 preconditioners = [jnp.eye(max_size) for s in shapes]
                 padded_statistics.extend(statistics)
                 padded_preconditioners.extend(preconditioners)
+                exponent = (
+                    preconditioner.exponent_for_preconditioner()
+                    if exponent_override == 0
+                    else exponent_override
+                )
+                exponents.extend([exponent] * len(shapes))
 
             diagonal_statistics = []
             if graft_type != GraftingType.SGD:
@@ -871,6 +951,7 @@ def distributed_shampoo(
                     _quantize_diagonal_statistics(diagonal_statistics),
                     _quantize_momentum(jnp.zeros_like(param)),
                     _quantize_momentum(jnp.zeros_like(param)),
+                    init_training_metrics(len(sizes)),
                     index_start,
                     sizes,
                 )
@@ -888,14 +969,238 @@ def distributed_shampoo(
         padded_preconditioners.extend(
             [jnp.eye(max_size, dtype=padded_statistics[0].dtype) for _ in range(to_pad)]
         )
+        exponents.extend([1 for _ in range(to_pad)])
         global_stats = GlobalShardedParameterStats(
-            jnp.stack(padded_statistics), jnp.stack(padded_preconditioners)
+            jnp.stack(padded_statistics),
+            jnp.stack(padded_preconditioners),
+            jnp.stack(exponents),
         )
         return ShampooState(
             count=jnp.zeros([], jnp.int32),
             stats=ShardedShampooStats(global_stats, local_stats),
         )
 
+    def _max_statistics_size_from_params(params):
+        max_size = 0
+        for param in params:
+            param_clone = jnp.zeros(param.shape, dtype=param.dtype)
+            preconditioner = Preconditioner(
+                param_clone, block_size, best_effort_shape_interpretation
+            )
+            if not _skip_preconditioning(param):
+                shapes = preconditioner.shapes_for_preconditioners()
+                sizes = [s[0] for s in shapes]
+                max_size = max(max(sizes), max_size)
+        return max_size
+
+    def _remove_leading_sharding_annotation(pspec):
+        """Mapping from N-d to (N-1)-d, used for quantization, factoring etc."""
+        # None and PSpec(None) are valid PSpecs.
+        if pspec and len(pspec) > 1:
+            return pjit.PartitionSpec(*pspec[1:])
+        else:
+            return None
+
+    def sharded_init_partition_spec_fn(
+        params, params_partition_spec, partition_spec_for_statistics
+    ):
+        """Returns a parallel state tree with PartitionSpec associated with state.
+
+
+        Args:
+          params: A pytree with params.
+          params_partition_spec: A pytree with PartitionSpec for params.
+          partition_spec_for_statistics: PartitionSpec for the statistics.
+        """
+        # Parallel lists of spec, and params.
+        param_pspec_flat, _ = jax.tree_flatten(
+            params_partition_spec, is_leaf=lambda x: x is None
+        )
+        params_flat, treedef = jax.tree_flatten(params)
+        assert param_pspec_flat
+        assert params_flat
+        # Step is replicated across cores.
+        # None means cores.
+        local_stats_flat = []
+        num_statistics = 0
+        for param, param_pspec in zip(params_flat, param_pspec_flat):
+            param_clone = jnp.zeros(param.shape, dtype=param.dtype)
+            preconditioner = Preconditioner(
+                param_clone, block_size, best_effort_shape_interpretation
+            )
+            shapes = preconditioner.shapes_for_preconditioners()
+            sizes = []
+
+            index_start = num_statistics
+            if not _skip_preconditioning(param):
+                sizes = [s[0] for s in shapes]
+                shapes = preconditioner.shapes_for_preconditioners()
+                num_statistics += len(shapes)
+
+            diagonal_statistics_pspec = []
+            diagonal_statistics_scale_pspec = []
+            if graft_type != GraftingType.SGD:
+                # Identically shaped param.
+                diagonal_statistics_pspec = param_pspec
+                if quantized_dtype_for_diagonal_statistics_buffers() != jnp.float32:
+                    diagonal_statistics_scale_pspec = (
+                        _remove_leading_sharding_annotation(param_pspec)
+                    )
+
+            m1_pspec = param_pspec
+            m2_pspec = param_pspec
+
+            m1_scale_pspec = []
+            m2_scale_pspec = []
+
+            if quantized_dtype_for_momentum_buffers() != jnp.float32:
+                m1_scale_pspec = _remove_leading_sharding_annotation(m1_pspec)
+                m2_scale_pspec = _remove_leading_sharding_annotation(m2_pspec)
+
+            local_stats_flat.append(
+                LocalShardedParameterStats(
+                    QuantizedValue(
+                        diagonal_statistics_pspec,
+                        [],
+                        diagonal_statistics_scale_pspec,
+                        quantized_dtype_for_diagonal_statistics_buffers(),
+                        False,
+                        list(param.shape),
+                    ),
+                    QuantizedValue(
+                        m1_pspec,
+                        [],
+                        m1_scale_pspec,
+                        quantized_dtype_for_momentum_buffers(),
+                        False,
+                        list(param.shape),
+                    ),
+                    QuantizedValue(
+                        m2_pspec,
+                        [],
+                        m2_scale_pspec,
+                        quantized_dtype_for_momentum_buffers(),
+                        False,
+                        list(param.shape),
+                    ),
+                    init_training_metrics_pspec(len(sizes)),
+                    index_start,
+                    sizes,
+                )
+            )
+
+        local_stats = jax.tree_unflatten(treedef, local_stats_flat)
+        global_stats = GlobalShardedParameterStats(
+            partition_spec_for_statistics,
+            partition_spec_for_statistics,
+            pjit.PartitionSpec(),
+        )
+        count_pspec = pjit.PartitionSpec()
+        return ShampooState(
+            count=count_pspec, stats=ShardedShampooStats(global_stats, local_stats)
+        )
+
+    def sharded_init_shape_and_dtype_fn(params):
+        """Returns a parallel state tree with shape, dtype associated with state.
+
+
+        Args:
+          params: A pytree with params.
+        """
+        # Parallel lists of spec, and params.
+        params_flat, treedef = jax.tree_flatten(params)
+        assert params_flat
+        # Step is replicated across cores.
+        # None means cores.
+        local_stats_flat = []
+        num_statistics = 0
+        for param in params_flat:
+            param_clone = jnp.zeros(param.shape, dtype=param.dtype)
+            preconditioner = Preconditioner(
+                param_clone, block_size, best_effort_shape_interpretation
+            )
+            shapes = preconditioner.shapes_for_preconditioners()
+            sizes = []
+
+            index_start = num_statistics
+            if not _skip_preconditioning(param):
+                sizes = [s[0] for s in shapes]
+                shapes = preconditioner.shapes_for_preconditioners()
+                num_statistics += len(shapes)
+
+            diagonal_statistics_shape_and_dtype = []
+            diagonal_statistics_scale_shape_and_dtype = []
+            if graft_type != GraftingType.SGD:
+                diagonal_statistics_shape_and_dtype = [list(param.shape), param.dtype]
+                qdtype = quantized_dtype_for_diagonal_statistics_buffers()
+                if qdtype != jnp.float32:
+                    diagonal_statistics_shape_and_dtype = [list(param.shape), qdtype]
+                    diagonal_statistics_scale_shape_and_dtype = [
+                        list(param.shape)[1:],
+                        param.dtype,
+                    ]
+
+            m1_shape_and_dtype = [list(param.shape), param.dtype]
+            m2_shape_and_dtype = [list(param.shape), param.dtype]
+
+            m1_scale_shape_and_dtype = []
+            m2_scale_shape_and_dtype = []
+
+            qdtype = quantized_dtype_for_momentum_buffers()
+            if qdtype != jnp.float32:
+                m1_shape_and_dtype = [list(param.shape), qdtype]
+                m2_shape_and_dtype = [list(param.shape), qdtype]
+
+                m1_scale_shape_and_dtype = [list(param.shape)[1:], qdtype]
+                m2_scale_shape_and_dtype = [list(param.shape)[1:], qdtype]
+
+            local_stats_flat.append(
+                LocalShardedParameterStats(
+                    QuantizedValue(
+                        diagonal_statistics_shape_and_dtype,
+                        [],
+                        diagonal_statistics_scale_shape_and_dtype,
+                        quantized_dtype_for_diagonal_statistics_buffers(),
+                        False,
+                        list(param.shape),
+                    ),
+                    QuantizedValue(
+                        m1_shape_and_dtype,
+                        [],
+                        m1_scale_shape_and_dtype,
+                        quantized_dtype_for_momentum_buffers(),
+                        False,
+                        list(param.shape),
+                    ),
+                    QuantizedValue(
+                        m2_shape_and_dtype,
+                        [],
+                        m2_scale_shape_and_dtype,
+                        quantized_dtype_for_momentum_buffers(),
+                        False,
+                        list(param.shape),
+                    ),
+                    init_training_metrics_shapes(len(sizes)),
+                    index_start,
+                    sizes,
+                )
+            )
+
+        local_stats = jax.tree_unflatten(treedef, local_stats_flat)
+        max_statistics_size = _max_statistics_size_from_params(params_flat)
+        to_pad = -num_statistics % num_devices_for_pjit
+        num_statistics += to_pad
+        statistics_shape = [num_statistics, max_statistics_size, max_statistics_size]
+        global_stats = GlobalShardedParameterStats(
+            [statistics_shape, jnp.float32],
+            [statistics_shape, jnp.float32],
+            [[num_statistics], jnp.int32],
+        )
+        return ShampooState(
+            count=[[], jnp.float32],
+            stats=ShardedShampooStats(global_stats, local_stats),
+        )
+
     def sharded_update_fn(grads, state, params):
         """Transform the input gradient and update all statistics in sharded mode.
 
@@ -923,20 +1228,6 @@ def distributed_shampoo(
             params_flat,
         )
 
-        exponents = []
-        for stat, param in zip(new_stats_flat, params_flat):
-            num_statistics = len(stat.statistics)
-            if num_statistics > 0:
-                preconditioner = Preconditioner(
-                    param, block_size, best_effort_shape_interpretation
-                )
-                exponent = (
-                    preconditioner.exponent_for_preconditioner()
-                    if exponent_override == 0
-                    else exponent_override
-                )
-                exponents.extend([exponent] * num_statistics)
-
         outputs = jax.tree_multimap(
             lambda g, s, p: _transform_grad(g, s, p, state.count),
             grads_flat,
@@ -951,7 +1242,6 @@ def distributed_shampoo(
             _convert_from_parameter_stats(new_stat, local_stat)
             for new_stat, local_stat in zip(new_stats_flat, local_stats_flat)
         ]
-        new_local_stats = jax.tree_unflatten(treedef, new_local_stats_flat)
 
         max_size = global_stats.statistics.shape[1]
         new_padded_statistics = []
@@ -974,22 +1264,16 @@ def distributed_shampoo(
                 for _ in range(to_pad)
             ]
         )
-        exponents.extend([1 for _ in range(to_pad)])
         new_stacked_padded_statistics = jnp.stack(new_padded_statistics)
-        new_stacked_exponents = jnp.stack(exponents)
-
-        def _matrix_inverse_pth_root_vmap(xs, ps):
-            mi_pth_root = functools.partial(
-                matrix_inverse_pth_root,
-                ridge_epsilon=matrix_epsilon,
-                precision=precision,
-            )
-            preconditioners, errors = jax.vmap(mi_pth_root)(xs, ps)
-            return preconditioners, errors
+        new_stacked_padded_statistics = pjit.with_sharding_constraint(
+            new_stacked_padded_statistics, statistics_partition_spec
+        )
 
         def _internal_inverse_pth_root_all():
-            preconditioners, errors = _matrix_inverse_pth_root_vmap(
-                new_stacked_padded_statistics, new_stacked_exponents
+            preconditioners, errors = _matrix_inverse_pth_root_pjit(
+                new_stacked_padded_statistics,
+                global_stats.exponents,
+                statistics_partition_spec,
             )
             return preconditioners, errors
 
@@ -1000,13 +1284,18 @@ def distributed_shampoo(
             # shaped tensors. Note statistics will be ignored as we are passing in
             # a large init value for error.
             preconditioners_init = new_stacked_padded_statistics
-            errors_init = np.stack([inverse_failure_threshold] * len(exponents))
+            n = new_stacked_padded_statistics.shape[0]
+            errors_init = jnp.ones([n], jnp.float32) * inverse_failure_threshold
             init_state = [preconditioners_init, errors_init]
             perform_step = state.count % preconditioning_compute_steps == 0
             new_preconditioners, errors = efficient_cond(
                 perform_step, _internal_inverse_pth_root_all, init_state
             )
 
+        new_local_stats_flat = _add_error_into_local_stats(
+            new_local_stats_flat, errors, inverse_failure_threshold
+        )
+        new_local_stats = jax.tree_unflatten(treedef, new_local_stats_flat)
         errors = errors.reshape((-1, 1, 1))
         predicate = jnp.logical_or(
             jnp.isnan(errors), errors >= inverse_failure_threshold
@@ -1017,7 +1306,9 @@ def distributed_shampoo(
             + (1.0 - predicate) * new_preconditioners
         )
         new_global_stats = GlobalShardedParameterStats(
-            new_stacked_padded_statistics, new_conditional_preconditioners
+            new_stacked_padded_statistics,
+            new_conditional_preconditioners,
+            global_stats.exponents,
         )
         new_shampoo_state = ShampooState(
             count=state.count + 1,
@@ -1048,6 +1339,7 @@ def distributed_shampoo(
                 _maybe_quantize_preconditioners(preconditioners),
                 _quantize_momentum(jnp.zeros_like(param)),
                 _quantize_momentum(jnp.zeros_like(param)),
+                init_training_metrics(len(statistics)),
             )
 
         return ShampooState(
@@ -1092,6 +1384,7 @@ def distributed_shampoo(
             state.preconditioners,
             state.diagonal_momentum,
             state.momentum,
+            state.training_metrics,
         )
 
     def _matrix_inverse_pth_root_vmap(xs, ps):
@@ -1115,33 +1408,27 @@ def distributed_shampoo(
 
         return jax.vmap(matrix_inverse_pth_root_wrapper)(qxs, qds, qbs, ps)
 
-    def _matrix_inverse_pth_root_pjit(xs, ps):
-        mesh_axis_names_tuple = tuple(mesh_axis_names)
+    def _matrix_inverse_pth_root_pjit(xs, ps, statistics_partition_spec=None):
         # Partition the concatenated statistics matrix across all cores.
-        partitioned_xs, partitioned_ps = pjit.pjit(
-            lambda x, y: (x, y),
-            in_axis_resources=None,
-            out_axis_resources=pjit.PartitionSpec(
-                mesh_axis_names_tuple,
-            ),
-        )(xs, ps)
+        pspec_for_partition = preconditioner_partition_spec
+        partitioned_xs = pjit.with_sharding_constraint(xs, pspec_for_partition)
+        partitioned_ps = pjit.with_sharding_constraint(
+            ps, pjit.PartitionSpec(preconditioner_partition_spec[0])
+        )
         # Run matrix inverse pth root on each shard.
         partitioned_preconditioners, partitioned_errors = _matrix_inverse_pth_root_vmap(
             partitioned_xs, partitioned_ps
         )
+        # Reshard output to have the same PSpec as input. This is required to avoid
+        # vmap seeing the full set of statistics.
+        partitioned_preconditioners = pjit.with_sharding_constraint(
+            partitioned_preconditioners, pspec_for_partition
+        )
         # Recombine the outputs at each core.
-        preconditioners, errors = pjit.pjit(
-            lambda x, y: (x, y),
-            in_axis_resources=(
-                pjit.PartitionSpec(
-                    mesh_axis_names_tuple,
-                ),
-                pjit.PartitionSpec(
-                    mesh_axis_names_tuple,
-                ),
-            ),
-            out_axis_resources=(None, None),
-        )(partitioned_preconditioners, partitioned_errors)
+        preconditioners = pjit.with_sharding_constraint(
+            partitioned_preconditioners, statistics_partition_spec
+        )
+        errors = pjit.with_sharding_constraint(partitioned_errors, pjit.PartitionSpec())
         return preconditioners, errors
 
     def _pmap_compute_preconditioners(
@@ -1223,31 +1510,54 @@ def distributed_shampoo(
             )
 
         new_preconditioners_flat = []
+        new_errors_flat = []
         for p, shape, prev_p, error in zip(
             preconditioners_flat, original_shapes, prev_preconditioners, errors_flat
         ):
             new_preconditioners_flat.append(
                 _select_preconditioner(error, p[: shape[0], : shape[1]], prev_p)
             )
+            new_errors_flat.append(error)
 
         assert len(states) == len(num_statistics_per_state)
         assert len(new_preconditioners_flat) == num_statistics
+        assert len(new_errors_flat) == num_statistics
 
         # Add back empty preconditioners so we that we can set the optimizer state.
         preconditioners_for_states = []
         idx = 0
+        errors_for_states = []
         for num_statistics, state in zip(num_statistics_per_state, states):
             if num_statistics == 0:
                 preconditioners_for_states.append([])
+                errors_for_states.append([])
             else:
                 preconditioners_for_state = new_preconditioners_flat[
                     idx : idx + num_statistics
                 ]
                 assert len(state.statistics) == len(preconditioners_for_state)
                 preconditioners_for_states.append(preconditioners_for_state)
+
+                errors_for_state = jnp.stack(
+                    new_errors_flat[idx : idx + num_statistics]
+                )
+                assert len(state.statistics) == len(errors_for_state)
+                errors_for_states.append(errors_for_state)
+
                 idx += num_statistics
         new_states = []
-        for state, new_preconditioners in zip(states, preconditioners_for_states):
+        for state, new_preconditioners, new_errors in zip(
+            states, preconditioners_for_states, errors_for_states
+        ):
+            if state.statistics:
+                new_errors = jnp.where(
+                    jnp.logical_and(
+                        new_errors > 0.0, new_errors != inverse_failure_threshold
+                    ),
+                    new_errors,
+                    state.training_metrics.inverse_pth_root_errors,
+                )
+            new_training_metrics = TrainingMetrics(new_errors)
             new_states.append(
                 ParameterStats(
                     state.diagonal_statistics,
@@ -1255,6 +1565,7 @@ def distributed_shampoo(
                     new_preconditioners,
                     state.diagonal_momentum,
                     state.momentum,
+                    new_training_metrics,
                 )
             )
 
@@ -1413,6 +1724,7 @@ def distributed_shampoo(
         new_quantized_preconditioners_flat = []
         new_quantized_diagonals_flat = []
         new_quantized_bucket_sizes_flat = []
+        new_errors_flat = []
         for p, d, b, shape, prev_p, error in zip(
             quantized_preconditioners_flat,
             quantized_diagonals_flat,
@@ -1432,6 +1744,7 @@ def distributed_shampoo(
             new_quantized_bucket_sizes_flat.append(
                 _select_preconditioner(error, b[: shape[0]], prev_p.bucket_size)
             )
+            new_errors_flat.append(error)
 
         assert len(states) == len(num_statistics_per_state)
         assert len(new_quantized_preconditioners_flat) == num_statistics
@@ -1440,10 +1753,12 @@ def distributed_shampoo(
 
         # Add back empty preconditioners so we that we can set the optimizer state.
         preconditioners_for_states = []
+        errors_for_states = []
         idx = 0
         for num_statistics, state in zip(num_statistics_per_state, states):
             if num_statistics == 0:
                 preconditioners_for_states.append([])
+                errors_for_states.append([])
             else:
                 quantized_preconditioners_for_state = (
                     new_quantized_preconditioners_flat[idx : idx + num_statistics]
@@ -1454,10 +1769,14 @@ def distributed_shampoo(
                 quantized_bucket_sizes_for_state = new_quantized_bucket_sizes_flat[
                     idx : idx + num_statistics
                 ]
+                errors_for_state = jnp.stack(
+                    new_errors_flat[idx : idx + num_statistics]
+                )
 
                 assert len(state.statistics) == len(quantized_preconditioners_for_state)
                 assert len(state.statistics) == len(quantized_diagonals_for_state)
                 assert len(state.statistics) == len(quantized_bucket_sizes_for_state)
+                assert len(state.statistics) == len(errors_for_state)
 
                 quantized_preconditioners = []
                 for qv, qd, qb in zip(
@@ -1469,9 +1788,21 @@ def distributed_shampoo(
                         QuantizedValue(qv, qd, qb, qv.dtype, True, list(qv.shape))
                     )
                 preconditioners_for_states.append(quantized_preconditioners)
+                errors_for_states.append(errors_for_state)
                 idx += num_statistics
         new_states = []
-        for state, new_preconditioners in zip(states, preconditioners_for_states):
+        for state, new_preconditioners, new_errors in zip(
+            states, preconditioners_for_states, errors_for_states
+        ):
+            if state.statistics:
+                new_errors = jnp.where(
+                    jnp.logical_and(
+                        new_errors > 0.0, new_errors != inverse_failure_threshold
+                    ),
+                    new_errors,
+                    state.training_metrics.inverse_pth_root_errors,
+                )
+            new_training_metrics = TrainingMetrics(new_errors)
             new_states.append(
                 ParameterStats(
                     state.diagonal_statistics,
@@ -1479,6 +1810,7 @@ def distributed_shampoo(
                     new_preconditioners,
                     state.diagonal_momentum,
                     state.momentum,
+                    new_training_metrics,
                 )
             )
 
@@ -1560,31 +1892,53 @@ def distributed_shampoo(
             )
 
         new_preconditioners_flat = []
+        new_errors_flat = []
         for p, shape, prev_p, error in zip(
             preconditioners_flat, original_shapes, prev_preconditioners, errors_flat
         ):
             new_preconditioners_flat.append(
                 _select_preconditioner(error, p[: shape[0], : shape[1]], prev_p)
             )
+            new_errors_flat.append(error)
 
         assert len(states) == len(num_statistics_per_state)
         assert len(new_preconditioners_flat) == num_statistics
 
         # Add back empty preconditioners so we that we can set the optimizer state.
         preconditioners_for_states = []
+        errors_for_states = []
         idx = 0
         for num_statistics, state in zip(num_statistics_per_state, states):
             if num_statistics == 0:
                 preconditioners_for_states.append([])
+                errors_for_states.append([])
             else:
                 preconditioners_for_state = new_preconditioners_flat[
                     idx : idx + num_statistics
                 ]
                 assert len(state.statistics) == len(preconditioners_for_state)
                 preconditioners_for_states.append(preconditioners_for_state)
+
+                errors_for_state = jnp.stack(
+                    new_errors_flat[idx : idx + num_statistics]
+                )
+                assert len(state.statistics) == len(errors_for_state)
+                errors_for_states.append(errors_for_state)
                 idx += num_statistics
+
         new_states = []
-        for state, new_preconditioners in zip(states, preconditioners_for_states):
+        for state, new_preconditioners, new_errors in zip(
+            states, preconditioners_for_states, errors_for_states
+        ):
+            if state.statistics:
+                new_errors = jnp.where(
+                    jnp.logical_and(
+                        new_errors > 0.0, new_errors != inverse_failure_threshold
+                    ),
+                    new_errors,
+                    state.training_metrics.inverse_pth_root_errors,
+                )
+            new_training_metrics = TrainingMetrics(new_errors)
             new_states.append(
                 ParameterStats(
                     state.diagonal_statistics,
@@ -1592,6 +1946,7 @@ def distributed_shampoo(
                     new_preconditioners,
                     state.diagonal_momentum,
                     state.momentum,
+                    new_training_metrics,
                 )
             )
 
@@ -1778,7 +2133,9 @@ def distributed_shampoo(
             state.preconditioners,
             _quantize_momentum(grafting_update_with_wd_momentum),
             _quantize_momentum(shampoo_update_with_wd_momentum),
+            state.training_metrics,
         )
+
         return transformed_update, param_stats
 
     def update_fn(grads, state, params):
@@ -1821,6 +2178,15 @@ def distributed_shampoo(
         return updates, new_state
 
     if shard_optimizer_states:
-        return optax.GradientTransformation(sharded_init_fn, sharded_update_fn)
+        # Hijacks the init_fn signature so we can return an OptState with
+        # appropriate init_fns.
+        def _init_fns(unused_params):
+            return InitFnState(
+                init_fn=sharded_init_fn,
+                pspec_fn=sharded_init_partition_spec_fn,
+                shape_and_dtype_fn=sharded_init_shape_and_dtype_fn,
+            )
+
+        return optax.GradientTransformation(_init_fns, sharded_update_fn)
     else:
         return optax.GradientTransformation(init_fn, update_fn)

tools/train/train.py

@@ -25,7 +25,7 @@ import sys
 import time
 from dataclasses import asdict, dataclass, field
 from pathlib import Path
-from typing import Callable, Optional
+from typing import Any, Callable, NamedTuple, Optional
 
 import datasets
 import jax
@@ -36,12 +36,12 @@ import transformers
 import wandb
 from datasets import Dataset
 from distributed_shampoo import GraftingType, distributed_shampoo
-from flax.core.frozen_dict import freeze
+from flax.core.frozen_dict import FrozenDict, freeze, unfreeze
 from flax.serialization import from_bytes, to_bytes
 from flax.training import train_state
 from flax.training.common_utils import onehot, stack_forest
 from jax.experimental import PartitionSpec, maps
-from jax.experimental.pjit import pjit
+from jax.experimental.pjit import pjit, with_sharding_constraint
 from tqdm import tqdm
 from transformers import HfArgumentParser
 
@@ -248,6 +248,10 @@ class TrainingArguments:
         default=1024,
         metadata={"help": "Chunked size for large layers with Distributed Shampoo."},
     )
+    start_preconditioning_step: int = field(
+        default=100,
+        metadata={"help": "Number of steps before starting to update preconditioner."},
+    )
     preconditioning_compute_steps: int = field(
         default=10, metadata={"help": "Number of steps to update preconditioner."}
     )
@@ -478,6 +482,7 @@ def main():
             artifact_dir,
             dtype=getattr(jnp, model_args.dtype),
             abstract_init=True,
+            load_on_cpu=True,
         )
 
         # load tokenizer
@@ -501,12 +506,14 @@ def main():
                 seed=training_args.seed_model,
                 dtype=getattr(jnp, model_args.dtype),
                 abstract_init=True,
+                load_on_cpu=True,
             )
         else:
             model = DalleBart(
                 config,
                 seed=training_args.seed_model,
                 dtype=getattr(jnp, model_args.dtype),
+                load_on_cpu=True,
             )
 
         # Load tokenizer
@@ -520,6 +527,12 @@ def main():
                 use_fast=True,
             )
 
+    # get PartitionSpec for model params (required to be a dict)
+    param_spec = set_partitions(model.params)
+
+    # convert params to frozen dict
+    model._params = freeze(model.params)
+
     # Preprocessing the datasets.
     # We need to normalize and tokenize inputs and targets.
 
@@ -536,14 +549,14 @@ def main():
 
     # Store some constant
     num_epochs = training_args.num_train_epochs
-    # batch size per node
-    train_batch_size = (
-        training_args.per_device_train_batch_size * jax.local_device_count()
+    # batch size
+    minibatch_size = (
+        training_args.per_device_train_batch_size * training_args.dp_devices
     )
-    batch_size_per_node = train_batch_size * training_args.gradient_accumulation_steps
+    batch_size_per_node = minibatch_size * training_args.gradient_accumulation_steps
     batch_size_per_step = batch_size_per_node * jax.process_count()
     eval_batch_size = (
-        training_args.per_device_eval_batch_size * jax.local_device_count()
+        training_args.per_device_eval_batch_size * training_args.dp_devices
     )
     len_train_dataset, len_eval_dataset = dataset.length
     steps_per_epoch = (
@@ -599,14 +612,17 @@ def main():
             beta2=training_args.beta2,
             diagonal_epsilon=1e-10,
             matrix_epsilon=1e-8,
-            start_preconditioning_step=training_args.warmup_steps,
+            start_preconditioning_step=training_args.start_preconditioning_step,
             preconditioning_compute_steps=training_args.preconditioning_compute_steps,
             statistics_compute_steps=1,
             best_effort_shape_interpretation=True,
             graft_type=GraftingType.RMSPROP_NORMALIZED,
             nesterov=False,
             exponent_override=0,
-            batch_axis_name="batch",
+            statistics_partition_spec=PartitionSpec(None, "batch", None),
+            preconditioner_partition_spec=PartitionSpec("batch", None, None),
+            num_devices_for_pjit=training_args.dp_devices,
+            shard_optimizer_states=True,
             inverse_failure_threshold=0.1,
             moving_average_for_momentum=True,
             skip_preconditioning_dim_size_gt=training_args.skip_preconditioning_dim_size_gt,
@@ -614,6 +630,13 @@ def main():
             precision=jax.lax.Precision.HIGHEST,
             best_effort_memory_usage_reduction=training_args.optim_quantized,
         )
+        # get the real optimizer and helper functions
+        update_fn = optimizer.update
+        optimizer = optimizer.init(model.params)
+        opt_fn = NamedTuple("opt_fn", pspec_fn=Any, shape_and_dtype_fn=Any)(
+            optimizer.pspec_fn, optimizer.shape_and_dtype_fn
+        )
+        optimizer = optax.GradientTransformation(optimizer.init_fn, update_fn)
 
     elif training_args.optim == "adam":
         optimizer = optax.adamw(
@@ -630,31 +653,45 @@ def main():
             clipping_threshold=training_args.max_grad_norm,
         )
 
-    # get opt_state shape without actual init
-    opt_state_shape = jax.eval_shape(lambda x: optimizer.init(x), model.params)
+    # get PartitionSpec for optimizer state
+    def get_opt_state_spec_and_shape(param_spec):
+        if training_args.optim in ["adam", "adafactor"]:
+            # get opt_state shape without actual init
+            opt_state_shape = jax.eval_shape(optimizer.init, model.params)
+
+            if training_args.optim == "adam":
+
+                def _opt_state_spec_per_leaf(x):
+                    if isinstance(x, FrozenDict):
+                        # variables with same structure as params
+                        return param_spec
+                    else:
+                        # other variables such as count
+                        return None
+
+                opt_state_spec = jax.tree_map(
+                    _opt_state_spec_per_leaf,
+                    opt_state_shape,
+                    # return None spec for empty elements
+                    is_leaf=lambda x: isinstance(x, (FrozenDict, optax.EmptyState)),
+                )
 
-    # get PartitionSpec for model params
-    param_spec = set_partitions(model.params)
+            elif training_args.optim == "adafactor":
+                # factorized state must be replicated (rank different than params)
+                opt_state_spec = None
 
-    # create PartitionSpec for opt_state
-    def opt_state_spec_per_leaf(x):
-        if training_args.optim in ["adam", "adafactor"]:
-            if isinstance(x, dict):
-                # variables with same structure as params
-                return param_spec
-            else:
-                # other variables such as count
-                return None
+        elif training_args.optim == "distributed_shampoo":
+            opt_state_spec = opt_fn.pspec_fn(
+                params=model.params,
+                params_partition_spec=param_spec,
+                partition_spec_for_statistics=PartitionSpec(None, "batch", None),
+            )
+            opt_state_shape = opt_fn.shape_and_dtype_fn(model.params)
         else:
-            # TODO: create spec for Distributed Shampoo
             raise NotImplementedError
+        return opt_state_spec, opt_state_shape
 
-    opt_state_spec = jax.tree_map(
-        opt_state_spec_per_leaf,
-        opt_state_shape,
-        # return None spec for empty elements
-        is_leaf=lambda x: isinstance(x, (dict, optax.EmptyState)),
-    )
+    opt_state_spec, opt_state_shape = get_opt_state_spec_and_shape(param_spec)
 
     # create a mesh
     mesh_shape = (training_args.dp_devices, training_args.mp_devices)
@@ -674,51 +711,61 @@ def main():
         tx=optimizer,
     )
 
-    opt_state, attr_state = None, None
-    if training_args.resume_from_checkpoint is not None:
-        # restore opt_state
-        with (Path(artifact_dir) / "opt_state.msgpack").open("rb") as f:
-            opt_state = from_bytes(opt_state_shape, f.read())
-            # need to freeze dict for pjit
-            opt_state = jax.tree_map(
-                lambda x: freeze(x) if isinstance(x, dict) else x,
-                opt_state,
-                is_leaf=lambda x: isinstance(x, (dict, optax.EmptyState)),
-            )
-        # restore other attributes
-        with (Path(artifact_dir) / "training_state.json").open("r") as f:
-            attr_state = json.load(f)
-
     # create training state
-    def init_state(params, opt_state):
+    with maps.mesh(mesh.devices, mesh.axis_names):
         if training_args.resume_from_checkpoint is None:
-            state = TrainState.create(
-                apply_fn=model.__call__,
-                tx=optimizer,
-                params=freeze(params),
-                dropout_rng=dropout_rng,
-            )
+
+            def init_state(params):
+                return TrainState.create(
+                    apply_fn=model.__call__,
+                    tx=optimizer,
+                    params=params,
+                    dropout_rng=dropout_rng,
+                )
+
+            state = pjit(
+                init_state,
+                in_axis_resources=(param_spec,),
+                out_axis_resources=state_spec,
+                donate_argnums=(0,),
+            )(model.params)
+
         else:
-            state = TrainState(
-                apply_fn=model.__call__,
-                tx=optimizer,
-                params=freeze(params),
-                opt_state=opt_state,
-                dropout_rng=dropout_rng,
-                **attr_state,
-            )
-        return state
+            # restore opt_state
+            with (Path(artifact_dir) / "opt_state.msgpack").open("rb") as f:
+                opt_state = from_bytes(opt_state_shape, f.read())
+
+            # restore other attributes
+            with (Path(artifact_dir) / "training_state.json").open("r") as f:
+                attr_state = json.load(f)
+
+            def restore_state(params, opt_state):
+                return TrainState(
+                    apply_fn=model.__call__,
+                    tx=optimizer,
+                    params=params,
+                    opt_state=opt_state,
+                    dropout_rng=dropout_rng,
+                    **attr_state,
+                )
 
-    with maps.mesh(mesh.devices, mesh.axis_names):
-        state = pjit(
-            init_state,
-            in_axis_resources=(param_spec, opt_state_spec),
-            out_axis_resources=state_spec,
-            donate_argnums=(0, 1),
-        )(freeze(model.params), opt_state)
+            state = pjit(
+                restore_state,
+                in_axis_resources=(param_spec, opt_state_spec),
+                out_axis_resources=state_spec,
+                donate_argnums=(0, 1),
+            )(model.params, opt_state)
+
+            # remove opt_state from CPU
+            del opt_state
 
-    # free memory from large parameters
-    del model._params, opt_state
+    # free memory
+    del model._params
+
+    # define batch specs
+    keys = ["attention_mask", "decoder_input_ids", "input_ids", "labels"]
+    batch_spec = freeze({k: PartitionSpec("batch") for k in keys})
+    grad_batch_spec = freeze({k: PartitionSpec(None, "batch") for k in keys})
 
     # label smoothed cross entropy
     def loss_fn(logits, labels):
@@ -728,11 +775,24 @@ def main():
 
     # Define gradient update step fn
     def train_step(state, batch, delta_time):
-        dropout_rng, new_dropout_rng = jax.random.split(state.dropout_rng)
-        # use a different rng per node
-        dropout_rng = jax.random.fold_in(dropout_rng, jax.process_index())
+        # batch is (gradient_accumulation_steps, minibatch_size, ...)
+        # check correct batch shape during compilation
+        assert batch["labels"].shape[0:3] == (
+            training_args.gradient_accumulation_steps,
+            training_args.dp_devices,
+            training_args.per_device_train_batch_size,
+        ), f"Expected label batch of shape dp_devices x gradient_acculumation x batch_per_device and got {batch['labels'].shape}"
+
+        # get a minibatch (one gradient accumulation slice)
+        def get_minibatch(batch, grad_idx):
+            return jax.tree_map(
+                lambda x: jax.lax.dynamic_index_in_dim(x, grad_idx, keepdims=False),
+                batch,
+            )
 
-        def compute_loss(params, minibatch):
+        def compute_loss(params, minibatch, dropout_rng):
+            # minibatch has dim (batch_size, ...)
+            minibatch = unfreeze(minibatch)
             labels = minibatch.pop("labels")
             logits = state.apply_fn(
                 **minibatch, params=params, dropout_rng=dropout_rng, train=True
@@ -741,36 +801,75 @@ def main():
 
         grad_fn = jax.value_and_grad(compute_loss)
 
+        def loss_and_grad(grad_idx, dropout_rng):
+            # minibatch at grad_idx, shape (dp_devices, per_device_train_batch_size, ...)
+            minibatch = get_minibatch(batch, grad_idx)
+            # ensure batch is sharded over devices
+            minibatch = jax.tree_map(
+                lambda x: with_sharding_constraint(x, PartitionSpec("batch")), minibatch
+            )
+            # calculate loss and grads independently per dp_device
+            loss_grads = jax.vmap(grad_fn, in_axes=(None, 0, None), out_axes=(0, 0))(
+                state.params, minibatch, dropout_rng
+            )
+            # ensure they are sharded over devices
+            loss_grads = jax.tree_map(
+                lambda x: with_sharding_constraint(x, PartitionSpec("batch")),
+                loss_grads,
+            )
+
+            # average across all devices
+            loss_grads = jax.tree_map(lambda x: jnp.mean(x, axis=0), loss_grads)
+
+            # return loss and grads
+            return loss_grads
+
+        # create a new rng
+        dropout_rng, _ = jax.random.split(state.dropout_rng)
+        # use a different rng per node
+        dropout_rng = jax.random.fold_in(dropout_rng, jax.process_index())
+
         if training_args.gradient_accumulation_steps == 1:
-            minibatch = jax.tree_map(lambda x: x[0], batch)
-            loss, grads = grad_fn(state.params, minibatch)
-        else:
 
-            def _cumul_loss_grads(i, cumul_loss_grads):
-                minibatch = jax.tree_map(lambda x: x[i], batch)
-                return jax.tree_map(
-                    lambda x, y: x + y,
-                    cumul_loss_grads,
-                    grad_fn(state.params, minibatch),
-                )
+            def batch_step(dropout_rng):
+                dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)
+                loss_grad = loss_and_grad(0, dropout_rng)
+                return loss_grad, new_dropout_rng
 
-            init_loss_grads = (
+            loss_grad, dropout_rng = batch_step(dropout_rng)
+        else:
+            # create initial state for per_minibatch_step loop
+            init_cumul_loss_grad = (
                 0.0,
                 jax.tree_map(jnp.zeros_like, state.params),
             )
-            loss, grads = jax.tree_map(
-                lambda x: x / training_args.gradient_accumulation_steps,
-                jax.lax.fori_loop(
-                    0,
-                    training_args.gradient_accumulation_steps,
-                    _cumul_loss_grads,
-                    init_loss_grads,
-                ),
+            init_minibatch_step = (init_cumul_loss_grad, dropout_rng)
+
+            # accumulate gradients
+            def cumul_minibatch_step(grad_idx, cumul_loss_grad_dropout):
+                cumul_loss_grad, dropout_rng = cumul_loss_grad_dropout
+                dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)
+                loss_grad = loss_and_grad(grad_idx, dropout_rng)
+                cumul_loss_grad = jax.tree_map(jnp.add, cumul_loss_grad, loss_grad)
+                return cumul_loss_grad, new_dropout_rng
+
+            # loop over gradients
+            loss_grad, dropout_rng = jax.lax.fori_loop(
+                0,
+                training_args.gradient_accumulation_steps,
+                cumul_minibatch_step,
+                init_minibatch_step,
+            )
+            # sum -> mean
+            loss_grad = jax.tree_map(
+                lambda x: x / training_args.gradient_accumulation_steps, loss_grad
             )
 
+        # update state
+        loss, grads = loss_grad
         state = state.apply_gradients(
             grads=grads,
-            dropout_rng=new_dropout_rng,
+            dropout_rng=dropout_rng,
             train_time=state.train_time + delta_time,
             train_samples=state.train_samples + batch_size_per_step,
         )
@@ -784,6 +883,7 @@ def main():
 
     # Define eval fn
     def eval_step(params, batch):
+        batch = unfreeze(batch)
         labels = batch.pop("labels")
         logits = model(**batch, params=params, train=False)[0]
         loss = loss_fn(logits, labels)
@@ -795,13 +895,13 @@ def main():
     # Create parallel version of the train and eval step
     p_train_step = pjit(
         train_step,
-        in_axis_resources=(state_spec, PartitionSpec("batch", None), None),
+        in_axis_resources=(state_spec, grad_batch_spec, None),
         out_axis_resources=(state_spec, None),
         donate_argnums=(0,),
     )
     p_eval_step = pjit(
         eval_step,
-        in_axis_resources=(param_spec, PartitionSpec("batch", None)),
+        in_axis_resources=(param_spec, batch_spec),
         out_axis_resources=None,
     )
 
@@ -842,9 +942,7 @@ def main():
         # ======================== Evaluating ==============================
         eval_metrics = []
         if training_args.do_eval:
-            eval_loader = dataset.dataloader(
-                "eval", training_args.per_device_eval_batch_size
-            )
+            eval_loader = dataset.dataloader("eval", eval_batch_size)
             eval_steps = (
                 len_eval_dataset // eval_batch_size
                 if len_eval_dataset is not None
@@ -857,8 +955,8 @@ def main():
                 leave=False,
                 total=eval_steps,
             ):
-                # Model forward
-                metrics = p_eval_step(state.params, batch)
+                # TODO: make this more efficient once training loop is fast
+                metrics = p_eval_step(state.params, freeze(batch))
                 eval_metrics.append(metrics)
 
             # normalize eval metrics
@@ -962,8 +1060,7 @@ def main():
             # Generate an epoch by shuffling sampling indices from the train dataset
             train_loader = dataset.dataloader(
                 "train",
-                training_args.per_device_train_batch_size,
-                training_args.gradient_accumulation_steps,
+                batch_size_per_node,
                 epoch,
             )
             # train
@@ -974,15 +1071,27 @@ def main():
                 leave=False,
                 total=steps_per_epoch,
             ):
-
                 # calculate delta time (we have a lag of one step but it's ok)
                 new_time = time.perf_counter()
                 delta_time = new_time - last_time
                 last_time = new_time
 
+                # reshape data into (gradient_accumulation_steps, dp_devices, batch_per_dp, ...)
+                batch = jax.tree_map(
+                    lambda x: x.reshape(
+                        (
+                            training_args.gradient_accumulation_steps,
+                            training_args.dp_devices,
+                            training_args.per_device_train_batch_size,
+                        )
+                        + x.shape[1:]
+                    ),
+                    batch,
+                )
+
                 # train step
-                state, train_metrics = p_train_step(state, batch, delta_time)
-                step = state.step
+                state, train_metrics = p_train_step(state, freeze(batch), delta_time)
+                step = int(state.step)
 
                 if step % training_args.logging_steps == 0 and jax.process_index() == 0:
                     all_metrics = metrics_logger.get_all_train_metrics(