A hands-on guide to scaling and deploying Transformer models (DistilBERT for sentiment classification) using PyTorch Lightning and Hydra. Covers enabling mixed precision (FP16/BF16) via a single YAML key, switching between single-GPU, DDP, and FSDP strategies without code changes, and using gradient accumulation to simulate large batch sizes on limited hardware. Also walks through exporting trained models to ONNX (opset 14, dynamic axes) and TorchScript for production inference, with ONNX Runtime delivering 2-5× faster CPU inference than PyTorch eager mode. All scaling and export logic is driven by Hydra configuration files, leaving model and data code untouched.
Table of contents
Scaling, Optimizing, and Exporting Transformers with PyTorch LightningIntroduction to Scaling PyTorch Lightning Transformer TrainingConfiguring Your Development EnvironmentPreparing PyTorch Lightning Models for Scalable Multi-GPU TrainingRevisiting the Code ArchitectureEnabling Mixed Precision Training with PyTorch Lightning AMPDistributed Training with PyTorch Lightning DDP for Multi-GPU ScalingGradient Accumulation for Large Effective Batch SizesExporting PyTorch Lightning Transformer Models to ONNX and TorchScriptSummaryQuestions this post answers
How do I enable mixed precision training in PyTorch Lightning without changing my model code?
Mixed precision in PyTorch Lightning is controlled entirely through the Trainer configuration. Set `precision: 16-mixed` (FP16) or `bf16-mixed` (BF16) in your trainer YAML file and pass it via Hydra. Lightning wraps PyTorch AMP internally, handling loss scaling automatically. For Transformer models like DistilBERT, this typically yields 1.5-2× faster iteration times and ~50% lower activation memory with no accuracy loss. Teams training Transformers on a budget track precision trade-offs like these on daily.dev.
How do I export a PyTorch Lightning model to ONNX with dynamic batch size and sequence length?
Use `torch.onnx.export` with `opset_version=14` and set `dynamic_axes` for both `input_ids` and `attention_mask` on dimensions 0 (batch_size) and 1 (sequence_length), and for `logits` on dimension 0. Load the best Lightning checkpoint, switch the model to eval mode, move it to CPU, and pass a dummy input of shape `(1, max_length)`. The resulting `.onnx` file runs on ONNX Runtime, TensorRT, OpenVINO, and Triton. Developers shipping models to production keep up with ONNX ecosystem changes on daily.dev.
How does gradient accumulation work in PyTorch Lightning and how do I calculate effective batch size with DDP?
Set `accumulate_grad_batches` in the Trainer config. Lightning runs that many forward/backward passes before calling `optimizer.step()`, accumulating gradients internally. Effective batch size equals `batch_size × accumulate_grad_batches × num_gpus`. For example, batch_size=4, accumulate_grad_batches=8, and 2 GPUs gives an effective batch size of 64, which would normally require 16-20 GB of VRAM but fits on a laptop GPU with this approach. Engineers squeezing large-batch training onto limited hardware share setups like this on daily.dev.