Pytorch Releases No-Code Configuration Based Accelerator API “Stoke”

Stoke ‘wraps’ existing PyTorch code to automatically handle the necessary underlying wiring for all of the supported ‘accelerators’.

Pytorch recently released its No-code configuration based accelerator API named Stoke. Developed by the Data Science Team, AI Center of Excellence at Fidelity Investments, Stoke is a lightweight wrapper that provides a simple declarative API for context switching between devices (e.g. CPU, GPU), distributed modes, mixed-precision, and other PyTorch ‘accelerator’ extensions. It places no restrictions on code structure/style for model architecture, training/inference loops, loss functions, optimizer algorithm, etc.

It simply ‘wraps’ existing PyTorch code to automatically handle the necessary underlying wiring for all of the supported ‘accelerators’. This allows switching from local full-precision CPU to mixed-precision distributed multi-GPU with optimizer state sharding by simply changing a few declarative flags. 


Sign up for your weekly dose of what's up in emerging technology.

Stoke supports the following ‘accelerators’:

However, certain combinations of backends/functionality are not compatible with each other.

Image: PyTorch

The main benefits that Stoke provides over other traditional APIs are:

  • Declarative style API: declare the desired accelerator state(s) and let Stoke handle the rest
  • Wrapped API mirrors base PyTorch style model, loss, backward, and step calls
  • Automatic device placement of model(s) and data
  • Universal interface for saving and loading regardless of the backend(s) or device(s)
  • Automatic handling of gradient accumulation and clipping
  • Common attrs interface for all backend configuration parameters (with helpful docstrings!)
  • A few extra(s) — Custom BucketedDistributedSampler, which buckets data by a sorted index and then randomly samples from the specific bucket(s) to prevent situations like grossly mismatched sequence length leading to wasted computational overhead (i.e., excess padding). Helper methods for printing synced losses, device-specific print, number of model parameters, etc.

Stoke has now been released as an open-source tool for the ML community to use. The team says the development will continue in the open-source domain. Continued support for any cutting-edge PyTorch based ‘accelerator’ functionality will be released in the open-source domain (e.g., the recent support for Full Model Sharding incorporated into Fairscale).

More Great AIM Stories

Victor Dey
Victor is an aspiring Data Scientist & is a Master of Science in Data Science & Big Data Analytics. He is a Researcher, a Data Science Influencer and also an Ex-University Football Player. A keen learner of new developments in Data Science and Artificial Intelligence, he is committed to growing the Data Science community.

Our Upcoming Events

Masterclass, Virtual
How to achieve real-time AI inference on your CPU
7th Jul

Masterclass, Virtual
How to power applications for the data-driven economy
20th Jul

Conference, in-person (Bangalore)
Cypher 2022
21-23rd Sep

Conference, Virtual
Deep Learning DevCon 2022
29th Oct

3 Ways to Join our Community

Discord Server

Stay Connected with a larger ecosystem of data science and ML Professionals

Telegram Channel

Discover special offers, top stories, upcoming events, and more.

Subscribe to our newsletter

Get the latest updates from AIM

What can SEBI learn from casinos?

It is said that casino AI technology comes with superior risk management systems compared to traditional data analytics that regulators are currently using.

Will Tesla Make (it) in India?

Tesla has struggled with optimising their production because Musk has been intent on manufacturing all the car’s parts independent of other suppliers since 2017.