Package: habana-lightning-plugins
Habana Lightning PluginsHabana Lightning plugins is a suite of plugins that aid/accelerate model training using Lightning framework for HPU.
The plugins acts as an extension to the lightning framework to support HPU specific features.Currently, the following plugins are available:HPUDataModuleHPUProfilerInstallationTo install Habana lightning plugins run the following command:python-umpipinstallhabana-lightning-pluginsHPUDataModuleHPUDataModuleis an extension to theLightningDataModuleclass which uses Habana's dataloader to load and pre-process the input data.
Using HPUDataModule offloads the data preprocessing overhead to the HPU and in turn increases the performance of training. The wrapper also
aids in switching between hardware and software preprocessor based on the specific Gaudi device used.VisitHabana Dataloaderfor more information related to Habana Dataloader.UsageThe following shows an example of how to use theHPUDataModule:Import Habana Datamodule:fromhabana_lightning_plugins.datamoduleimportHPUDataModuleCreate and initialize HPUDataModule object with the dataset and the configuration required to preprocess the data:train_dir="./path/to/train/data"val_dir="./path/to/val/data"normalize=transforms.Normalize(mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225])train_transforms=[transforms.RandomResizedCrop(224),transforms.RandomHorizontalFlip(),transforms.ToTensor(),normalize,]val_transforms=[transforms.Resize(256),transforms.CenterCrop(224),transforms.ToTensor(),normalize,]data_module=HPUDataModule(train_dir,val_dir,train_transforms=train_transforms,val_transforms=val_transforms,num_workers=8,batch_size=32,shuffle=False,pin_memory=True,drop_last=True,)Create an object of Lightning trainer and model:trainer=pl.Trainer(devices=1,accelerator="hpu",max_epochs=1,max_steps=2)model=RN50Module()# Or any other model to be defined by userPass the datamodule object as an argument to trainer to execute train/val/test loops:trainer.fit(model,datamodule=data_module)trainer.validate(model,datamodule=data_module)ExamplesA sample script can be found atexamples/hpu_datamodule_sample.py.pythonexamples/hpu_datamodule_sample.py--data-path<pathtoImagenetdataset-ILSVRC2012>A reference model using HPUDataModule can be found in theResNet50 Model ReferenceLimitationsHPUDataModule supports theImagenetdataset only.HPUDataModule supports only 8 parallel data loader workers.HPUProfilerHPUProfiler is a lightning implementation of PyTorch profiler for HPU devices. It aids in obtaining profiling summary of PyTorch functions.
It subclasses PyTorch Lightning'sPyTorch profiler.Default ProfilingFor auto profiling, create a HPUProfiler instance and pass it to trainer.
At the end ofprofiler.fit(), it will generate a json trace for the run.
In caseaccelerator = "hpu"is not used with HPUProfiler, then it will dump only CPU traces, similar to PyTorchProfiler.# Import profilerfromhabana_lightning_plugins.profilerimportHPUProfiler# Create profiler objectprofiler=HPUProfiler()accelerator="hpu"# Pass profiler to the trainertrainer=Trainer(profiler=profiler,accelerator=accelerator,)Distributed ProfilingTo profile a distributed model, use the HPUProfiler with the filename argument which will save a report per rank:fromhabana_lightning_plugins.profilerimportHPUProfilerprofiler=HPUProfiler(filename="perf-logs")trainer=Trainer(profiler=profiler,accelerator="hpu")Custom ProfilingToprofile custom actions of interest, reference a profiler in the LightningModule:fromhabana_lightning_plugins.profilerimportHPUProfiler# Reference profiler in LightningModuleclassMyModel(LightningModule):def__init__(self,profiler=None):self.profiler=profiler# To profile in any part of your code, use the self.profiler.profile() functiondefcustom_processing_step_basic(self,data):withself.profiler.profile("my_custom_action"):...returndata# Alternatively, use self.profiler.start("my_custom_action")# and self.profiler.stop("my_custom_action") functions# to enclose the part of code to be profiled.defcustom_processing_step_granular(self,data):self.profiler.start("my_custom_action")...self.profiler.stop("my_custom_action")returndata# Pass profiler instance to LightningModuleprofiler=HPUProfiler()model=MyModel(profiler)trainer=Trainer(profiler=profiler,accelerator="hpu")For more details on profiler, refer toPyTorchProfilerVisualize Profiled OperationsProfiler will dump traces in json format. The traces can be visualized in 2 ways:Using PyTorch TensorBoard ProfilerFor further instructions see,https://github.com/pytorch/kineto/tree/master/tb_plugin.# Install tensorboardpython-umpipinstalltensorboardtorch-tb-profiler# Start the TensorBoard server (default at port 6006):tensorboard--logdir./tensorboard--port6006# Now open the following url on your browserhttp://localhost:6006/#profileUsing Chrome1. Open Chrome and copy/paste this URL: `chrome://tracing/`.
2. Once tracing opens, click on `Load` at the top-right and load one of the generated traces.LimitationsWhen using the HPUProfiler, wall clock time will not be representative of the true wall clock time. This is due to forcing profiled operations to be measured synchronously, when many HPU ops happen asynchronously.
It is recommended to use this Profiler to find bottlenecks/breakdowns, however for end to end wall clock time use the SimpleProfiler.HPUProfiler.summary() is not supportedPassing profiler name as string "hpu" to the trainer is not supported.Supported ConfigurationsValidated onSynapseAI VersionPyTorch VersionPyTorch Lightning VersionGaudi1.9.01.13.11.9.4Gaudi21.9.01.13.11.9.4Changeloghabana-lightning-plugins introduced with support for datamodule and profiler pluginsContributingPull requests are welcome. For major changes, please open an issue first
to discuss what you would like to change.Please make sure to update tests as appropriate.