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

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

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. 

Stoke supports the following ‘accelerators’:

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

See Also
How Is Alteryx Advancing Data & Analytics Potential With Its Game-changing Resources

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).

What Do You Think?

Subscribe to our Newsletter

Get the latest updates and relevant offers by sharing your email.
Join our Telegram Group. Be part of an engaging community

Copyright Analytics India Magazine Pvt Ltd

Scroll To Top