When dealing with massive Transformer models, I need extreme speed. Is PyTorch Lightning the fastest way to train deep learning models, or should I go directly to something like Microsoft’s DeepSpeed? I've seen that Lightning has a DeepSpeed strategy, but is it just as fast as using the library natively?
3 answers
The beauty of the current ecosystem is that you don't have to choose. Lightning acts as an orchestrator for DeepSpeed. When you set strategy="deepspeed", it passes your model and configuration to the DeepSpeed engine. In my experience with billion-parameter models, there was zero performance difference between using native DeepSpeed and the Lightning wrapper. The advantage of Lightning here is that you can switch between different strategies (like FSDP or DDP) just by changing a single string in your code, which makes it the fastest way to experiment with different scaling techniques.
Do you think the integration handles the ZeRO-3 offloading as effectively as a manual configuration script would?
I found that Lightning’s built-in 16-bit mixed precision (AMP) is much more stable and easier to implement, which prevents training crashes and saves time.
Agreed. The fact that it handles the scaling of losses automatically in the background makes it far more reliable for long-running jobs.
Actually, Ronald, Lightning allows you to pass a full DeepSpeed config file or a dictionary. So, you get the exact same level of granular control over ZeRO-3 offloading, partition sizes, and overlap of communication. It just removes the need to manually wrap the optimizer and model. I’ve used it for several large-scale LLM finetuning tasks and the memory management was identical to the native implementation.