With the PyTorch 1.12 release hot off the press, let’s dive into the release notes and check out three exciting features in this release. For context, at PathAI, we use PyTorch pretty much exclusively for our Machine Learning infrastructure. The pythonic and robust nature of PyTorch makes it a natural fit for our workflows. You can check out our recently published case study on how we use PyTorch to improve patient outcomes.
Accelerated PyTorch Training on Mac
With the advent of the new M-series Apple processor, it seems like Moore’s Law is unlikely to find a local minimum anytime soon. While Apple pushes the new processor across the entire line of Macs, more and more developers have been pleading with tool makers to use all the new computational power; Tensorflow added support for Apple's M1 chips in Nov 2020 and it was only a matter of time before PyTorch released support as well. An issue opened by users of PyTorch requesting GPU accelerated training garnered input from over a thousand users on GitHub!
PyTorch core developer and project lead Soumith Chintala went on to say in November of 2021 that the team was working on it.
Fast forward almost eight months, and the feature is finally here! Those using PyTorch on Mac no longer need to sit while their CPU overworks itself trying to do backpropagation; you can use the new M1 Mac GPUs for accelerating PyTorch training! This is a huge advancement for the community, and it’s great to see PyTorch listen to the community suggestions.
You can read the full PyTorch M1-Mac announcement post on the PyTorch Blog.
Fully Sharded Data Parallel (FSDP) API
Sharding in ML is the practice of breaking up large monolithic data sets so they can be accessed and used more efficiently. If you aren’t familiar with Sharding, you might want to read this article from Meta’s Production Engineering Team which introduces the concept from the perspective of web development, but the same principles apply in the case of ML.
The PyTorch team recently published an article highlighting how the new Sharding API has been tested with models that contain over 1 trillion parameters. For context, many modern day machine learning models tend to have somewhere around 100 million parameters (though this number is going up significantly on a yearly basis). These techniques have been shown to help reduce the memory required for training by ~60% while still enabling larger models to be trained.
The new Beta FSDP API comes with different sharding techniques as well as faster model initialization which are sure to be big hits with the PyTorch developer community.
TorchArrow Beta Release
TorchArrow is a new package in the PyTorch ecosystem which is designed as a “torch.Tensor-like Python DataFrame library for data preprocessing in deep learning”. The logical question one might ask is: “I already use Pandas on a daily basis for my ML and Data Science workflows, why would I want to use TorchArrow?”. The answer, as you might see if you go through the “TorchArrow in 10 minutes” tutorial, is that TorchArrow is designed to just work with tensors. While Pandas can display some of the information about the tensor, it does not contain all of the details you might expect to see if you were using Pandas for the first time.
TorchArrow is built in top of Arrow (a project co-created by Wes Mckinney, who also made Pandas):
"Apache Arrow defines a language-independent columnar memory format for flat and hierarchical data, organized for efficient analytic operations on modern hardware like CPUs and GPUs. The Arrow memory format also supports zero-copy reads for lightning-fast data access without serialization overhead."
Common operations like map, filter, and reduce are all implemented in TorchArrow and follow similar syntactical paradigms as other PyTorch syntax which means you can work to have your entire ML workflow using a single toolkit rather than relying on other projects. TorchArrow will also include support for GPU accelerated workflows in a future release. Due to the Pandas-like API, the learning curve for TorchArrow should be limited. You can get started using it by following the installation instructions on GitHub.
PyTorch is one of the most popular machine learning frameworks with no signs of slowing down. This release demonstrates both efforts to improve the efficiency for existing developers while also ensuring that new use cases (specifically around topics like ML Ops) are being identified and designed for. At PathAI, we are working hard to ensure we leverage the latest developments in both the PyTorch ecosystem as well as the broader ML field.
PathAI is at the intersection of advanced machine learning and transforming pathology. If this interests you, then you may want to consider joining our growing team of machine learning engineers: https://www.pathai.com/join-our-team/careers/