Adding pad_to_multiple_of=16

run_mlm_flax_stream.py

@@ -218,9 +218,9 @@ class FlaxDataCollatorForLanguageModeling:
                 "You should pass `mlm=False` to train on causal language modeling instead."
             )
 
-    def __call__(self, examples: List[Dict[str, np.ndarray]]) -> Dict[str, np.ndarray]:
+    def __call__(self, examples: List[Dict[str, np.ndarray]], pad_to_multiple_of: int) -> Dict[str, np.ndarray]:
         # Handle dict or lists with proper padding and conversion to tensor.
-        batch = self.tokenizer.pad(examples, return_tensors=TensorType.NUMPY)
+        batch = self.tokenizer.pad(examples, pad_to_multiple_of=pad_to_multiple_of, return_tensors=TensorType.NUMPY)
 
         # If special token mask has been preprocessed, pop it from the dict.
         special_tokens_mask = batch.pop("special_tokens_mask", None)
@@ -653,7 +653,7 @@ if __name__ == "__main__":
             samples = advance_iter_and_group_samples(training_iter, train_batch_size, max_seq_length)
 
         # process input samples
-        model_inputs = data_collator(samples)
+        model_inputs = data_collator(samples, pad_to_multiple_of=16)
 
         # Model forward
         model_inputs = shard(model_inputs.data)
@@ -678,7 +678,7 @@ if __name__ == "__main__":
             for i, batch_idx in enumerate(tqdm(eval_batch_idx, desc="Evaluating ...", position=1)):
                 # process input samples
                 batch_eval_samples = {k: [v[idx] for idx in batch_idx] for k, v in eval_samples.items()}
-                model_inputs = data_collator(batch_eval_samples)
+                model_inputs = data_collator(batch_eval_samples, pad_to_multiple_of=16)
 
                 # Model forward
                 model_inputs = shard(model_inputs.data)