ipex-llm/docs/readthedocs/source/doc/Nano/Overview/pytorch_train.md

PyTorch Training

BigDL-Nano can be used to accelerate PyTorch or PyTorch-Lightning applications on training workloads. The optimizations in BigDL-Nano are delivered through an extended version of PyTorch-Lightning Trainer. These optimizations are either enabled by default or can be easily turned on by setting a parameter or calling a method.

We will briefly describe here the major features in BigDL-Nano for PyTorch training. You can find complete examples here.

Best Known Configurations

When you run source bigdl-nano-init, BigDL-Nano will export a few environment variables, such as OMP_NUM_THREADS and KMP_AFFINITY, according to your current hardware. Empirically, these environment variables work best for most PyTorch applications. After setting these environment variables, you can just run your applications as usual (python app.py) and no additional changes are required.

BigDL-Nano PyTorch Trainer

The PyTorch Trainer (bigdl.nano.pytorch.Trainer) is the place where we integrate most optimizations. It extends PyTorch Lightning's Trainer and has a few more parameters and methods specific to BigDL-Nano. The Trainer can be directly used to train a LightningModule.

For example,

from pytorch_lightning import LightningModule
from bigdl.nano.pytorch import Trainer

class MyModule(LightningModule):
    #  LightningModule definition

from bigdl.nano.pytorch import Trainer
lightning_module = MyModule()
trainer = Trainer(max_epoch=10)
trainer.fit(lightning_module, train_loader)

For regular PyTorch modules, we also provide a "compile" method, that takes in a PyTorch module, an optimizer, and other PyTorch objects and "compiles" them into a LightningModule.

For example,

from bigdl.nano.pytorch import Trainer
lightning_module = Trainer.compile(pytorch_module, loss, optimizer, scheduler)
trainer = Trainer(max_epoch=10)
trainer.fit(lightning_module, train_loader)

Intel® Extension for PyTorch

Intel Extension for PyTorch (a.k.a. IPEX) extends PyTorch with optimizations for an extra performance boost on Intel hardware. BigDL-Nano integrates IPEX through the Trainer. Users can turn on IPEX by setting use_ipex=True.

from bigdl.nano.pytorch import Trainer
trainer = Trainer(max_epoch=10, use_ipex=True)

Multi-instance Training

When training on a server with dozens of CPU cores, it is often beneficial to use multiple training instances in a data-parallel fashion to make full use of the CPU cores. However, using PyTorch's DDP API is a little cumbersome and error-prone, and if not configured correctly, it will make the training even slow.

BigDL-Nano makes it very easy to conduct multi-instance training. You can just set the num_processes parameter in the Trainer constructor and BigDL-Nano will launch the specific number of processes to perform data-parallel training. Each process will be automatically pinned to a different subset of CPU cores to avoid conflict and maximize training throughput.

from bigdl.nano.pytorch import Trainer
trainer = Trainer(max_epoch=10, num_processes=4)

Note that the effective batch size in multi-instance training is the batch_size in your dataloader times num_processes so the number of iterations of each epoch will be reduced num_processes fold. A common practice to compensate for that is to gradually increase the learning rate to num_processes times. You can find more details of this trick in this paper published by Facebook.

BigDL-Nano PyTorch TorchNano

The TorchNano (bigdl.nano.pytorch.TorchNano) class is what we use to accelerate raw pytorch code. By using it, we only need to make very few changes to accelerate custom training loop. For example,

from bigdl.nano.pytorch import TorchNano

class MyNano(TorchNano) :
    def train(self, ...):
        # copy your train loop here and make a few changes

MyNano().train(...)

• note: see this tutorial for details about our TorchNano.

Our TorchNano also integrates IPEX and distributed training optimizations. For example,

from bigdl.nano.pytorch import TorchNano

class MyNano(TorchNano):
    def train(self, ...):
        # define train loop

# enable IPEX optimizaiton
MyNano(use_ipex=True).train(...)

# enable IPEX and distributed training, using subprocess strategy
MyNano(use_ipex=True, num_processes=2, strategy="subprocess").train(...)

Optimized Data Pipeline

Computer Vision task often needs a data processing pipeline that sometimes constitutes a non-trivial part of the whole training pipeline. Leveraging OpenCV and libjpeg-turbo, BigDL-Nano can accelerate computer vision data pipelines by providing a drop-in replacement of torch_vision's datasets and transforms.

from bigdl.nano.pytorch.vision.datasets import ImageFolder
from bigdl.nano.pytorch.vision import transforms

data_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.ColorJitter(),
    transforms.RandomCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.Resize(128),
    transforms.ToTensor()
])

train_set = ImageFolder(train_path, data_transform)
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
trainer.fit(module, train_loader)