Handling Custom Habana Ops for PyTorch

Habana provides its own implementation of some complex PyTorch ops customized for Habana devices. In a given model, replacing these complex ops with custom Habana versions will give better performance.

Custom Optimizers

The following is a list of custom optimizers currently supported on Habana devices:

• FusedAdagrad - refer to torch.optim.Adagrad

• FusedAdamW - refer to AdamW from HuggingFace

• FusedLamb - refer to LAMB optimizer paper

• FusedSGD - refer to torch.optim.SGD

• Functional FusedAdamW - Functional version of the above FusedAdamW, based on torch.distributed.optim._FunctionalAdamW

Note

Functional FusedAdamW is available as habana_frameworks.torch.hpex.optimizers.distributed.FusedAdamW.`

The below shows an example code snippet demonstrating the usage of a custom optimizer:

try:
   from habana_frameworks.torch.hpex.optimizers import FusedLamb
except ImportError:
   raise ImportError("Please install habana_torch package")
   optimizer = FusedLamb(model.parameters(), lr=args.learning_rate)

Note

For models using Lazy mode execution, mark_step() must be added right after loss.backward() and optimizer.step().

Other Custom OPs

The following is a list of other custom ops currently supported on Habana devices:

• FusedClipNorm - refer to toch.nn.utils.clip_grad_norm_

The below is an example code snippet demonstrating the usage of FusedClipNorm:

try:
   from habana_frameworks.torch.hpex.normalization import FusedClipNorm
except ImportError:
   raise ImportError("Please install habana_torch package")
   FusedNorm = FusedClipNorm(model.parameters(), args.max_grad_norm)

FusedNorm.clip_norm(model.parameters())