fix(train): consider correct batch size

src/dalle_mini/data.py

@@ -156,21 +156,19 @@ class Dataset:
         self, split, per_device_batch_size, gradient_accumulation_steps=None, epoch=None
     ):
         num_devices = jax.local_device_count()
+        total_batch_size = per_device_batch_size * num_devices
+        if gradient_accumulation_steps is not None:
+            total_batch_size *= gradient_accumulation_steps
 
         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
+            steps_per_epoch = len(dataset) // total_batch_size
 
             if rng is not None:
                 batch_idx = jax.random.permutation(rng, len(dataset))
@@ -178,25 +176,24 @@ class Dataset:
                 batch_idx = jnp.arange(len(dataset))
 
             batch_idx = batch_idx[
-                : steps_per_epoch * batch_size
+                : steps_per_epoch * total_batch_size
             ]  # Skip incomplete batch.
-            batch_idx = batch_idx.reshape((steps_per_epoch, batch_size))
+            batch_idx = batch_idx.reshape((steps_per_epoch, total_batch_size))
 
             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:]),
+                        lambda x: x.reshape(
+                            (gradient_accumulation_steps, -1) + 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"]
@@ -214,19 +211,13 @@ class Dataset:
                 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]]) == total_batch_size:
                         batch = {k: jnp.array(v) for k, v in batch.items()}
                         if gradient_accumulation_steps is not None:
+                            # training mode
                             batch = jax.tree_map(
                                 lambda x: x.reshape(
-                                    (-1, per_device_batch_size) + x.shape[1:]
+                                    (gradient_accumulation_steps, -1) + x.shape[1:]
                                 ),
                                 batch,
                             )
@@ -242,15 +233,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):

tools/train/train.py

@@ -549,11 +549,11 @@ def main():
 
     # Store some constant
     num_epochs = training_args.num_train_epochs
-    # batch size per node
-    train_batch_size = (
+    # batch size
+    minibatch_size = (
         training_args.per_device_train_batch_size * jax.local_device_count()
     )
-    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()
@@ -770,6 +770,12 @@ def main():
 
     # Define gradient update step fn
     def train_step(state, batch, delta_time):
+        # check correct batch shape during compilation
+        assert batch["labels"].shape[0:2] == (
+            training_args.gradient_accumulation_steps,
+            minibatch_size,
+        ), f"Expected label batch of shape gradient_acculumation x minibatch_size x items and got {batch['labels'].shape}"
+        # create a new rng
         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())
@@ -837,13 +843,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, PartitionSpec(None, "batch"), 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, PartitionSpec("batch")),
         out_axis_resources=None,
     )