Porting PyTorch Models to Gaudi

This section details the steps required to prepare PyTorch models to run on Gaudi. To create and execute a simple PyTorch model with Gaudi, see Getting Started with PyTorch and Gaudi.

  1. Import Habana Torch Library:

import habana_frameworks.torch.core as htcore
  1. Target the Gaudi HPU device:

device = torch.device("hpu")
  1. Add mark_step(). In Lazy mode, mark_step() must be added in all training scripts right after loss.backward() and optimizer.step().



  • Placing mark_step() at any arbitrary point in the code is currently not supported.

  • Migrating BERT-like models to run on Gaudi when module weights are shared among two or more layers requires these weights to be shared after moving the model to the HPU device. For more details, refer to Weight Sharing.