feat: handle model parallel

src/dalle_mini/data.py

@@ -85,7 +85,12 @@ class Dataset:
                     else self.eval_dataset.select(range(self.max_eval_samples))
                 )
 
-    def preprocess(self, tokenizer, decoder_start_token_id, normalize_text, max_length):
+    def preprocess(self, tokenizer, config):
+        # get required config variables
+        decoder_start_token_id = config.decoder_start_token_id
+        normalize_text = config.normalize_text
+        max_length = config.max_text_length
+
         if self.streaming:
             # we need to shuffle early in streaming mode
             if hasattr(self, "train_dataset"):

src/dalle_mini/model/configuration.py

@@ -59,6 +59,7 @@ class DalleBartConfig(PretrainedFromWandbMixin, PretrainedConfig):
         is_encoder_decoder=True,
         forced_eos_token_id=None,
         tie_word_embeddings=False,  # different modalities and sizes
+        do_sample=True,
         **kwargs,
     ):
         self.normalize_text = normalize_text
@@ -87,28 +88,28 @@ class DalleBartConfig(PretrainedFromWandbMixin, PretrainedConfig):
             scale_embedding  # scale factor will be sqrt(d_model) if True
         )
 
-        # remove inferred keys to prevent errors when loading config (passed as kwargs)
-        for k in [
-            "pad_token_id",
-            "bos_token_id",
-            "eos_token_id",
-            "decoder_start_token_id",
-            "min_length",
-            "max_length",
-        ]:
-            kwargs.pop(k, None)
+        # special token id's are appended to vocab if not provided
+        decoder_start_token_id = kwargs.pop("decoder_start_token_id", image_vocab_size)
+        bos_token_id = kwargs.pop("bos_token_id", image_vocab_size)
+        pad_token_id = kwargs.pop("pad_token_id", image_vocab_size)
+        eos_token_id = kwargs.pop("eos_token_id", image_vocab_size)
+
+        # we generate to image_length + 1 (for bos) by default
+        min_length = kwargs.pop("min_length", image_length + 1)
+        max_length = kwargs.pop("max_length", image_length + 1)
 
         super().__init__(
-            pad_token_id=image_vocab_size
-            + 1,  # needed to avoid errors during generation (converted to jnp.array)
-            bos_token_id=image_vocab_size + 1,  # set to unreachable values
-            eos_token_id=image_vocab_size + 1,
+            # args required in parent class
             is_encoder_decoder=is_encoder_decoder,
-            decoder_start_token_id=image_vocab_size,  # BOS appended to vocab
-            forced_eos_token_id=forced_eos_token_id,
             tie_word_embeddings=tie_word_embeddings,
-            min_length=image_length + 1,
-            max_length=image_length + 1,
+            forced_eos_token_id=forced_eos_token_id,
+            decoder_start_token_id=decoder_start_token_id,
+            bos_token_id=bos_token_id,
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            min_length=min_length,
+            max_length=max_length,
+            do_sample=do_sample,
             **kwargs,
         )
 

src/dalle_mini/model/modeling.py

@@ -54,7 +54,7 @@ logger = logging.get_logger(__name__)
 class FlaxBartAttention(FlaxBartAttention):
     """
     Edits:
-    - causal mask is used only in decoder and considers image_length + 1 (for BOS)
+    - causal mask is used only in decoder and considers image_length
     """
 
     def setup(self) -> None:
@@ -81,7 +81,7 @@ class FlaxBartAttention(FlaxBartAttention):
         if self.causal:
             # used only in decoder
             self.causal_mask = make_causal_mask(
-                jnp.ones((1, self.config.image_length + 1), dtype="bool"), dtype="bool"
+                jnp.ones((1, self.config.image_length), dtype="bool"), dtype="bool"
             )
 
 
@@ -240,7 +240,7 @@ class FlaxBartDecoder(FlaxBartDecoder):
     """
     Edits:
     - offset set to 0 (no padding token)
-    - use image_length + 1 (for BOS) instead of max_position_embeddings
+    - use image_length instead of max_position_embeddings
     - use custom FlaxBartDecoderLayerCollection
     - embed_tokens cannot be None (issue at compile time)
     """
@@ -258,7 +258,7 @@ class FlaxBartDecoder(FlaxBartDecoder):
         # and adjust num_embeddings appropriately. Other models don't have this hack
         self.offset = 0
         self.embed_positions = nn.Embed(
-            self.config.image_length + 1 + self.offset,  # image length + 1 for BOS
+            self.config.image_length + self.offset,  # image length for BOS
             embed_dim,
             embedding_init=jax.nn.initializers.normal(self.config.init_std),
         )

tools/train/train.py

@@ -99,7 +99,7 @@ class ModelArguments:
 
     def __post_init__(self):
         if self.restore_state:
-            assert (
+            assert self.model_name_or_path is not None and (
                 "/model-" in self.model_name_or_path
             ), "Restoring state only available with W&B artifact reference"
             self.state_artifact = self.model_name_or_path.replace(
@@ -222,12 +222,13 @@ class TrainingArguments:
     )
 
     per_device_train_batch_size: int = field(
-        default=8, metadata={"help": "Batch size per GPU/TPU/CPU for training."}
+        default=8,
+        metadata={"help": "Batch size per data parallel device for training."},
     )
     per_device_eval_batch_size: Optional[int] = field(
         default=None,
         metadata={
-            "help": "Batch size per GPU/TPU/CPU for evaluation. Same as training batch size if not set."
+            "help": "Batch size per data parallel device for evaluation. Same as training batch size if not set."
         },
     )
 
@@ -523,12 +524,7 @@ def main():
     # Preprocessing the datasets.
     # We need to normalize and tokenize inputs and targets.
 
-    dataset.preprocess(
-        tokenizer=tokenizer,
-        decoder_start_token_id=model.config.decoder_start_token_id,
-        normalize_text=model.config.normalize_text,
-        max_length=model.config.max_text_length,
-    )
+    dataset.preprocess(tokenizer=tokenizer, config=model.config)
 
     # Initialize our training
     rng = jax.random.PRNGKey(training_args.seed_model)
@@ -874,9 +870,17 @@ def main():
 
     # Define eval fn
     def eval_step(state, batch):
-        batch, labels = batch.pop("labels")
-        logits = model(**batch, params=state.params, train=False)[0]
-        loss = loss_fn(logits, labels)
+        def compute_eval_loss(batch):
+            batch, labels = batch.pop("labels")
+            logits = state.apply_fn(**batch, params=state.params, train=False)[0]
+            return loss_fn(logits, labels)
+
+        # calculate loss independently per dp_device
+        loss = jax.vmap(compute_eval_loss, in_axes=(0,), out_axes=0)(batch)
+        # ensure they are sharded over dp devices
+        loss = with_sharding_constraint(loss, PartitionSpec("batch"))
+        # average across all devices
+        loss = jnp.mean(loss)
         return loss
 
     # Create parallel version of the train and eval step
@@ -946,7 +950,18 @@ def main():
                 leave=False,
                 total=eval_steps,
             ):
-                # freeze batch to pass safely to JAX transforms
+                # reshape data into (dp_devices, batch_per_dp, ...)
+                batch = jax.tree_map(
+                    lambda x: x.reshape(
+                        (
+                            training_args.dp_devices,
+                            training_args.per_device_eval_batch_size,
+                        )
+                        + x.shape[1:]
+                    ),
+                    batch,
+                )
+                # freeze batch to pass safely to jax transforms
                 batch = freeze(batch)
                 # accumulate losses async
                 eval_loss.append(p_eval_step(state, batch))