Quickstart
🤗 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 training 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(
model=model,
args=training_args,
train_dataset=train_dataset
... # other arguments
)where gaudi_config_name is the name of a model from the Hub (Gaudi configurations are stored in model repositories) or a path to a local Gaudi configuration file (you can see here how to write your own).
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 = "runwayml/stable-diffusion-v1-5"
- scheduler = DDIMScheduler.from_pretrained(model_name, subfolder="scheduler")
+ scheduler = GaudiDDIMScheduler.from_pretrained(model_name, subfolder="scheduler")
- pipeline = StableDiffusionPipeline.from_pretrained(
+ pipeline = GaudiStableDiffusionPipeline.from_pretrained(
model_name,
scheduler=scheduler,
+ use_habana=True,
+ use_hpu_graphs=True,
+ gaudi_config="Habana/stable-diffusion",
)
outputs = pipeline(
["An image of a squirrel in Picasso style"],
num_images_per_prompt=16,
+ batch_size=4,
)Ready-to-Use Examples
Here are examples for various modalities and tasks that can be used out of the box: