Why You Should Upgrade Your Code to PyTorch 2.0

Introduction

torch.set_default_device: This enables to setup your device globally! no more to("cuda") calls all over your code. Just set it up once and you're done 😎. This new feature enables you to remove some boiler plate on your code––note that this is optional and you can keep your code as it is if you'd prefer.
And the torch.compile function that promises extreme throughput on NVIDIA GPUs. Simply put, you wrap your model with compile_model = torch.compile(model) and you are good to go! Normally, you get a nice performance boost from 10% to 100%, depending on the model! 🚀

Benchmarking PyTorch 2.0

Model throughput (sample per second)
Model + DataLoader throughput (samples per sec)
Total runtime (secs)

from time import perf_counter

for epoch in range(epochs):
  tf = perf_counter()
  for x,y in dataloader:
    t0 = perf_counter()
    # pytorch train steps...
    x,y = x.to("cuda"), y.to("cuda")
    out = model(x)
    loss = loss_func(out, y)
    loss.backward()
    # other stuff like optimizer, schedulers, grad scalers...
    tf_with_dataloader = perf_counter() - tf # time with data gathering
    tf = perf_counter()
    wandb.log({"samples_per_sec":len(x)/(tf-t0),
               "sampler_per_sec_dl":len(x)/tf_with_dataloader})

ResNet50

The samples per second for the forward + backward pass (without the time spent gathering the data from the DataLoader)
The samples per second, including the data gathering.

BERT Training

CPU utilization of the DataLoader

Conclusions

The performance gain for the ResNet50 is around 33% on an A100 and around 25% on a V100
For the BERT model, the gain is around 45% on an A100 and 50% on a V100.