30,31d29
< from datasets import load_dataset
< 
32a31
> from datasets import load_dataset
44,46d42
<     Seq2SeqTrainer,
<     Seq2SeqTrainingArguments,
<     default_data_collator,
52a49,52
> from optimum.graphcore import IPUConfig, IPUSeq2SeqTrainer
> from optimum.graphcore import IPUSeq2SeqTrainingArguments as Seq2SeqTrainingArguments
> from optimum.graphcore.utils import check_min_version as gc_check_min_version
> 
56a57,59
> # Will error if the minimal version of Optimum Graphcore is not installed. Remove at your own risks.
> gc_check_min_version("0.6.0.dev0")
> 
146c149
<         default=128,
---
>         default=1024,
286,290d288
<     # Log on each process the small summary:
<     logger.warning(
<         f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
<         + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
<     )
370a369,374
>     ipu_config = IPUConfig.from_pretrained(
>         training_args.ipu_config_name if training_args.ipu_config_name else model_args.model_name_or_path,
>         cache_dir=model_args.cache_dir,
>         revision=model_args.model_revision,
>         use_auth_token=True if model_args.use_auth_token else None,
>     )
523,531c527,534
<     if data_args.pad_to_max_length:
<         data_collator = default_data_collator
<     else:
<         data_collator = DataCollatorForSeq2Seq(
<             tokenizer,
<             model=model,
<             label_pad_token_id=label_pad_token_id,
<             pad_to_multiple_of=8 if training_args.fp16 else None,
<         )
---
>     data_collator = DataCollatorForSeq2Seq(
>         tokenizer,
>         model=model,
>         # If data_args.pad_to_max_length then already padded during preprocessing.
>         padding=False if data_args.pad_to_max_length else "max_length",
>         label_pad_token_id=label_pad_token_id,
>         pad_to_multiple_of=None,
>     )
564c567
<     trainer = Seq2SeqTrainer(
---
>     trainer = IPUSeq2SeqTrainer(
565a569
>         ipu_config=ipu_config,