Handling Custom Habana Ops for PyTorch
On this Page
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())