Optimum documentation


Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started


🤗 Optimum Graphcore was designed with one goal in mind: making training and evaluation straightforward for any 🤗 Transformers user while leveraging the complete power of IPUs.

There are two main classes one needs to know:

  • IPUTrainer: the trainer class that takes care of compiling the model to run on IPUs, and of performing training and evaluation, check here for more information..
  • IPUConfig: the class that specifies attributes and configuration parameters to compile and put the model on the device, check here for more information.

The IPUTrainer is very similar to the 🤗 Transformers Trainer, and adapting a script using the Trainer to make it work with IPUs will mostly consists of simply swapping the Trainer class for the IPUTrainer one. That’s how most of the example scripts were adapted from their original counterparts.

-from transformers import Trainer, TrainingArguments
+from optimum.graphcore import IPUConfig, IPUTrainer, IPUTrainingArguments

-training_args = TrainingArguments(
+training_args = IPUTrainingArguments(
+    ipu_config_name="Graphcore/bert-base-ipu",  # Or any other IPUConfig on the Hub or stored locally
+# Loading the IPUConfig needed by the IPUTrainer to compile and train the model on IPUs
+ipu_config = IPUConfig.from_pretrained(
+    training_args.ipu_config_name if training_args.ipu_config_name else model_args.model_name_or_path,

 # Initialize our Trainer
-trainer = Trainer(
+trainer = IPUTrainer(
+    ipu_config=ipu_config,
     train_dataset=train_dataset if training_args.do_train else None,
     ...  # Other arguments