Optimum documentation


Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started


🤗 Optimum Habana was designed with one goal in mind: making training and evaluation straightforward for any 🤗 Transformers user while leveraging the complete power of Gaudi processors. There are two main classes one needs to know:

  • GaudiTrainer: the trainer class that takes care of compiling (lazy or eager mode) and distributing the model to run on HPUs, and of performing traning and evaluation.
  • GaudiConfig: the class that enables to configure Habana Mixed Precision and to decide whether optimized operators and optimizers should be used or not.

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

-from transformers import Trainer, TrainingArguments
+from optimum.habana import GaudiTrainer, GaudiTrainingArguments

# define the training arguments
-training_args = TrainingArguments(
+training_args = GaudiTrainingArguments(
+  use_habana=True,
+  use_lazy_mode=True,
+  gaudi_config_name=gaudi_config_name,

# Initialize our Trainer
-trainer = Trainer(
+trainer = GaudiTrainer(
    ... # other arguments

where gaudi_config_name is the name of a model from the Hub (Gaudi configurations are stored in model repositories). You can also give the path to a custom Gaudi configuration written in a JSON file such as this one:

  "use_habana_mixed_precision": true,
  "hmp_opt_level": "O1",
  "hmp_is_verbose": false,
  "use_fused_adam": true,
  "use_fused_clip_norm": true,
  "hmp_bf16_ops": [
  "hmp_fp32_ops": [

If you prefer to instantiate a Gaudi configuration to work on it before giving it to the trainer, you can do it as follows:

gaudi_config = GaudiConfig.from_pretrained(
    use_auth_token=True if model_args.use_auth_token else None,

Stable Diffusion

🤗 Optimum Habana also features HPU-optimized support for the 🤗 Diffusers library. Thus, you can easily deploy Stable Diffusion on Gaudi for performing text-to-image generation.

Here is how to use it and the differences with the 🤗 Diffusers library:

- from diffusers import DDIMScheduler, StableDiffusionPipeline
+ from optimum.habana.diffusers import GaudiDDIMScheduler, GaudiStableDiffusionPipeline

model_name = "CompVis/stable-diffusion-v1-4"

- scheduler = DDIMScheduler.from_pretrained(model_name, subfolder="scheduler")
+ scheduler = GaudiDDIMScheduler.from_pretrained(model_name, subfolder="scheduler")

- pipeline = StableDiffusionPipeline.from_pretrained(
+ pipeline = GaudiStableDiffusionPipeline.from_pretrained(
+   use_habana=True,
+   use_hpu_graphs=True,
+   gaudi_config="Habana/stable-diffusion",

outputs = generator(
    ["An image of a squirrel in Picasso style"],
+   batch_size=4,