feat: custom gradient accumulation

src/dalle_mini/data.py

@@ -153,16 +153,21 @@ class Dataset:
                     ),
                 )
 
-    def dataloader(self, split, batch_size, epoch=None):
+    def dataloader(
+        self, split, per_device_batch_size, gradient_accumulation_steps=None, epoch=None
+    ):
+        num_devices = jax.local_device_count()
+
         def _dataloader_datasets_non_streaming(
             dataset: Dataset,
-            batch_size: int,
+            per_device_batch_size: 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
             steps_per_epoch = len(dataset) // batch_size
 
             if rng is not None:
@@ -182,7 +187,11 @@ class Dataset:
                 yield batch
 
         def _dataloader_datasets_streaming(
-            dataset: Dataset, split: str, batch_size: int, epoch: int
+            dataset: Dataset,
+            split: str,
+            per_device_batch_size: int,
+            gradient_accumulation_steps: int,
+            epoch: int,
         ):
             keys = ["input_ids", "attention_mask", "labels", "decoder_input_ids"]
             batch = {k: [] for k in keys}
@@ -199,8 +208,22 @@ class Dataset:
                 for item in dataset:
                     for k, v in item.items():
                         batch[k].append(v)
-                    if len(batch[keys[0]]) == batch_size:
+                        # 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
+                    ):
                         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,
+                            )
                         batch = shard(batch)
                         yield batch
                         batch = {k: [] for k in keys}
@@ -214,11 +237,15 @@ class Dataset:
             raise ValueError(f'split must be "train" or "eval", got {split}')
 
         if self.streaming:
-            return _dataloader_datasets_streaming(ds, split, batch_size, epoch)
+            return _dataloader_datasets_streaming(
+                ds, split, per_device_batch_size, gradient_accumulation_steps, epoch
+            )
         else:
             if split == "train":
                 self.rng_dataset, input_rng = jax.random.split(self.rng_dataset)
-            return _dataloader_datasets_non_streaming(ds, batch_size, input_rng)
+            return _dataloader_datasets_non_streaming(
+                ds, per_device_batch_size, input_rng
+            )
 
     @property
     def length(self):

tools/train/train.py

@@ -277,8 +277,8 @@ class TrainingArguments:
         },
     )
 
-    num_train_epochs: float = field(
-        default=3.0, metadata={"help": "Total number of training epochs to perform."}
+    num_train_epochs: int = field(
+        default=3, metadata={"help": "Total number of training epochs to perform."}
     )
     warmup_steps: int = field(
         default=0, metadata={"help": "Linear warmup over warmup_steps."}
@@ -515,10 +515,10 @@ def main():
     rng, dropout_rng = jax.random.split(rng)
 
     # Store some constant
-    num_epochs = int(training_args.num_train_epochs)
+    num_epochs = training_args.num_train_epochs
     # batch size per node
     train_batch_size = (
-        int(training_args.per_device_train_batch_size) * jax.local_device_count()
+        training_args.per_device_train_batch_size * jax.local_device_count()
     )
     batch_size_per_update = (
         train_batch_size
@@ -526,7 +526,7 @@ def main():
         * jax.process_count()
     )
     eval_batch_size = (
-        int(training_args.per_device_eval_batch_size) * jax.local_device_count()
+        training_args.per_device_eval_batch_size * jax.local_device_count()
     )
     len_train_dataset, len_eval_dataset = dataset.length
     steps_per_epoch = (
@@ -645,12 +645,6 @@ def main():
             clipping_threshold=training_args.max_grad_norm,
         )
 
-    # add gradient accumulation
-    if training_args.gradient_accumulation_steps > 1:
-        optimizer = optax.MultiSteps(
-            optimizer, training_args.gradient_accumulation_steps
-        )
-
     # Setup train state
     state = TrainState.create(
         apply_fn=model.__call__,
@@ -673,16 +667,42 @@ def main():
     def train_step(state, batch, delta_time):
         dropout_rng, new_dropout_rng = jax.random.split(state.dropout_rng)
 
-        def compute_loss(params, batch):
-            labels = batch.pop("labels")
+        def compute_loss(params, minibatch):
+            labels = minibatch.pop("labels")
             logits = state.apply_fn(
-                **batch, params=params, dropout_rng=dropout_rng, train=True
+                **minibatch, params=params, dropout_rng=dropout_rng, train=True
             )[0]
-            loss = loss_fn(logits, labels)
-            return loss
+            return loss_fn(logits, labels)
 
         grad_fn = jax.value_and_grad(compute_loss)
-        loss, grads = grad_fn(state.params, batch)
+
+        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),
+                )
+
+            init_loss_grads = (
+                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,
+                ),
+            )
+
         grads = jax.lax.pmean(grads, "batch")
         state = state.apply_gradients(
             grads=grads,
@@ -871,7 +891,12 @@ def main():
         metrics_logger.log({"train/epoch": epoch}, step=unreplicate(state.step))
 
         # Generate an epoch by shuffling sampling indices from the train dataset
-        train_loader = dataset.dataloader("train", train_batch_size, epoch)
+        train_loader = dataset.dataloader(
+            "train",
+            training_args.per_device_train_batch_size,
+            training_args.gradient_accumulation_steps,
+            epoch,
+        )
         # train
         for batch in tqdm(
             train_loader,