Artificial intelligence development startup Lightning AI Inc. today announced the general availability of PyTorch Lightning 2.0, an open-source AI framework that’s used by thousands of organizations to train and scale up machine learning models.

PyTorch Lightning is one of the most popular machine learning development tools, with an average of 4 million downloads per month, thanks to the way it simplifies many manual tasks involved in creating AI-powered software. Building an AI application means not only creating a neural network, but also setting up and managing the infrastructure it runs on. Doing so involves lots of manual work and is very time-consuming, so developers can save time by using PyTorch Lightning to automate much of it.

Another task that PyTorch Lightning simplifies is AI training. It does so using a combination of large numbers of processors and tools that make it easier for developers to evaluate the speed of their neural networks and troubleshoot any issues that come up. Another benefit of PyTorch Lightning is that it reduces the amount of custom code developers must write to use multiple types of chips, such as graphics processing units and central processing units, in their ML training projects.

Lightning AI sells a commercial version of PyTorch Lightning that makes it simpler for enterprises to deploy and use. It also provides prepackaged AI applications for common use cases, which can be used by customers as blueprints for their own, custom applications.

With the launch of PyTorch Lightning 2.0, Lightning AI said, it’s introducing a stable application programming interface, more powerful features with a smaller footprint, and tools that make it easier to read and debug. It’s also introducing Lightning Fabric, which is a new library that gives developers more control over the AI training process, with features like callbacks and checkpoints and support for reinforcement learning, active learning and transformers.

According to Lightning AI, customers that would rather use a simple, scalable training method that works out of the box can rely on PyTorch Lightning 2.0, while those that need more granularity and control should use Lightning Fabric.

Lightning AI says this is a compelling choice because previously developers have been forced to go to extremes: Either they use proscriptive tools for training and deploying machine learning models, or they figure out everything themselves. Now, they get an extensive choice of training options.

If users require features such as checkpointing, logging, and early stopping to be integrated out of the box, they can use the original PyTorch Lightning framework. However, if they need to scale their models to multiple GPUs and write their own training loops for tasks such as reinforcement learning, the same method used by OpenAI LLC’s ChatGPT, then they should use Lightning Fabric.

Lightning AI co-founder and Chief Executive William Falcon said machine learning developers can now choose between additional granularity or a simple scaling solution that works out of the box.

“As the size of models deployed in enterprise settings continues to expand, the ability to scale the training process in a way that is simple, lightweight and cost-effective becomes increasingly important,” he said. “Lightning AI is proud to be expanding the choices that users have to scale their models for downstream tasks whether those are internal, like analytics and content creation tools, or customer-facing, like ChatGPT and other customer service chatbots.”

Images: Lightning AI