From PyTorch to PyTorch Fabric
Scale model training with 4 small changes.
AI isn’t magic. It’s math.
Understand the concepts powering technology like ChatGPT in minutes a day with Brilliant.
Thousands of quick, interactive lessons in AI, programming, logic, data science, and more make it easy. Try it free for 30 days.
Thanks to Brilliant for sponsoring today’s issue.
PyTorch gives so much flexibility and control. But it leads to a ton of boilerplate code.
PyTorch Lightning, however:
Massively reduces the boilerplate code.
Allow us to use distributed training features like DDP, FSDP, DeepSpeed, mixed precision training, and more, by directly specifying parameters:
But it isn’t as flexible as PyTorch to write manual training loops and stuff.
Lately, I have been experimenting with Lightning Fabric, which brings together:
The flexibility of PyTorch.
And direct access to distributed training features like PyTorch Lightning does.
Today’s issue covers the 4 small changes you can make to your existing PyTorch code to easily scale it to the largest billion-parameter models/LLMs.
All this is summarized in a single frame below:
Let’s distill it!
PyTorch to Lightning Fabric
Begin by importing the lightning module (install using pip install lightning), creating a Fabric object and launching it.
In my opinion, Fabric object is the most powerful aspect here since you can specify the training configuration and settings directly with parameters.
For instance, with a single line of code, you specify the number of devices and the parallelism strategy to use:
Moreover, it lets you take full advantage of the hardware on your system. It supports CPU, TPU, and GPU (NVIDIA, AMD, Apple Silicon).
Moving on, configure the model, the optimizer, and the dataloader as follows:
Next, remove all .to() and .cuda() calls since Fabric takes care of it automatically:
Finally, replace the loss.backward() call by fabric.backward(loss) call.
Done!
Now, you can train the model as you usually would.
That was simple, wasn’t it?
Here’s the documentation if you want to dive into more details about the usage: Lightning Fabric.
👉 Over to you: What are some issues with PyTorch?
P.S. For those wanting to develop “Industry ML” expertise:
At the end of the day, all businesses care about impact. That’s it!
Can you reduce costs?
Drive revenue?
Can you scale ML models?
Predict trends before they happen?
We have discussed several other topics (with implementations) in the past that align with such topics.
Here are some of them:
Learn sophisticated graph architectures and how to train them on graph data: A Crash Course on Graph Neural Networks – Part 1.
So many real-world NLP systems rely on pairwise context scoring. Learn scalable approaches here: Bi-encoders and Cross-encoders for Sentence Pair Similarity Scoring – Part 1.
Learn techniques to run large models on small devices: Quantization: Optimize ML Models to Run Them on Tiny Hardware.
Learn how to generate prediction intervals or sets with strong statistical guarantees for increasing trust: Conformal Predictions: Build Confidence in Your ML Model’s Predictions.
Learn how to identify causal relationships and answer business questions: A Crash Course on Causality – Part 1
Learn how to scale ML model training: A Practical Guide to Scaling ML Model Training.
Learn techniques to reliably roll out new models in production: 5 Must-Know Ways to Test ML Models in Production (Implementation Included)
Learn how to build privacy-first ML systems: Federated Learning: A Critical Step Towards Privacy-Preserving Machine Learning.
Learn how to compress ML models and reduce costs: Model Compression: A Critical Step Towards Efficient Machine Learning.
All these resources will help you cultivate key skills that businesses and companies care about the most.
SPONSOR US
Get your product in front of 110k+ data scientists and other tech professionals.
Our newsletter puts your products and services directly in front of an audience that matters — thousands of leaders, senior data scientists, machine learning engineers, data analysts, etc., who have influence over significant tech decisions and big purchases.
To ensure your product reaches this influential audience, reserve your space here or reply to this email to ensure your product reaches this influential audience.