October 2023 Summaries

A PyTorch vision transformer fine-tuning example demonstrates how Lightning Fabric can add advanced training options with minimal code changes while preserving the baseline’s performance of 94.06% test accuracy, 16.88 minutes of training time, and 16.70 GB of GPU memory on a single A100 GPU. Enabling 16-bit mixed precision through Fabric reduces training time by more than 30% and memory use by over 25% without reducing accuracy, because it combines fast low-precision computations with 32-bit operations for stability. Full FP16 training can produce NaN losses due to its limited numerical range, whereas full bfloat16 training offers a wider range and, in these experiments, maintained accuracy while roughly halving both training time and peak memory consumption relative to standard 32-bit training. The techniques are presented as broadly applicable to other architectures, although performance and accuracy should be verified for each model and bfloat16 requires compatible hardware such as NVIDIA A100-class GPUs.

PyTorch Lightning is presented as a high-level PyTorch framework that reduces training boilerplate and supports experimentation, scalable multi-device training, and large-model workflows. Its central components are the LightningModule, which organizes model initialization, training, validation, testing, prediction, metrics, and optimizer configuration, and the Trainer, which automates device management, mixed precision, distributed training, and training configuration. The walkthrough demonstrates building a CIFAR-10 image classifier by loading and normalizing data with torchvision, defining a compact convolutional neural network, wrapping it in a LightningModule with cross-entropy loss, accuracy metrics, and an SGD optimizer, and training it for up to 30 epochs. It also describes using TensorBoard to visualize logged loss and accuracy metrics and callbacks such as EarlyStopping and ModelCheckpoint to monitor and control training, while noting that LightningDataModule can further organize datasets and dataloaders for distributed environments.

The tutorial demonstrates how to train a small Transformer-based language model on the WikiText2 dataset using PyTorch Lightning, emphasizing that Lightning extends rather than replaces PyTorch by organizing existing PyTorch models into reproducible training workflows. It introduces Lightning’s core components: LightningDataModule for dataset preparation, splitting, and DataLoaders; LightningModule for wrapping the Transformer, defining forward passes, training and evaluation steps, loss logging, and optimization; and Trainer for managing the training loop and runtime configuration. The example uses a custom data module to automatically retrieve and split WikiText2, a custom model module using negative log-likelihood loss and SGD optimization, and configurable Trainer features including automatic device selection, checkpointing, early stopping, Weights & Biases logging, and PyTorch profiling. It concludes that this structure separates data, model, engineering, and support concerns, helping researchers and developers build more maintainable, scalable, and reusable deep-learning projects while retaining the flexibility of PyTorch.

PyTorch Lightning is an open-source wrapper for PyTorch that reduces boilerplate in deep learning projects by organizing model logic, training steps, optimization, logging, and evaluation through standardized components such as LightningModule and Trainer. It aims to make code cleaner, more modular, reproducible, and easier to maintain while allowing researchers to focus on model architecture rather than manual training-loop implementation. Migrating from traditional PyTorch generally involves moving forward-pass, loss, and optimizer logic into a LightningModule, then using a Trainer to run training. The framework also integrates with experiment-tracking tools such as TensorBoard and Weights & Biases, and provides built-in support for multi-GPU, distributed, mixed-precision, and lower-precision training with minimal code changes.

An experimental study of LoRA instruction fine-tuning used the Alpaca dataset, Llama 2 7B, and several language, factual, truthfulness, and arithmetic benchmarks to examine optimizer, quantization, training-duration, and adapter-hyperparameter choices on a single A100 GPU. QLoRA reduced peak GPU memory from about 21 GB to 14 GB but increased training time by roughly 30% and caused modest performance trade-offs, while replacing AdamW with SGD yielded little memory benefit at low LoRA ranks because relatively few adapter parameters were trainable. A cosine learning-rate scheduler improved some truthfulness and factual metrics, but training for two passes over the dataset generally degraded results, particularly arithmetic, likely because Alpaca contains limited arithmetic content. Expanding LoRA from attention components to all applicable layers and increasing rank improved results only when the alpha scaling parameter was adjusted proportionally; the strongest reported configuration used rank 256 and alpha 512, although very large ranks performed worse or failed to converge. The final leaderboard comparison found that tuned QLoRA improved several benchmarks over the base model but not all, underscoring that dataset selection, alignment methods, and evaluation diversity remain important alongside LoRA settings.

PyTorch Lightning’s Fully Sharded Data Parallel strategy is presented as a way to train and fine-tune billion-parameter models such as Llama, Stable Diffusion, and genomic models while addressing GPU memory limits, distributed-training complexity, and infrastructure costs. The approach shards model parameters, gradients, and optimizer states across GPUs, with optional CPU offloading, while Lightning’s Trainer and LightningModule provide configurable distributed execution, precision handling, callbacks, and customization hooks. Using a 2-billion-parameter Transformer example, the blog reports that memory usage falls from 39.02 GB on one GPU to 29.30 GB per GPU with default four-GPU FSDP, then to about 13 GB through layer-specific FSDP wrapping and 12.5 GB with activation checkpointing, which trades memory savings for slower computation. It then applies these techniques to Llama 7B using Lit-GPT, FSDP wrapping and checkpointing for transformer blocks, bfloat16 precision, and configurable gradient accumulation and validation settings, while noting that CPU offloading can further reduce GPU memory but may substantially slow training.

Quantization reduces machine-learning model memory requirements by representing parameters with lower-precision formats such as 8-bit or 4-bit values rather than float32 or float16, enabling large models such as Llama 2 and SDXL to run on memory-constrained edge and consumer hardware. Because it compresses values through specialized representations rather than merely removing bits, quantization can affect model accuracy and should be evaluated in its deployed form. Lightning Fabric supports bitsandbytes-based quantization through its `BitsandbytesPrecision` plugin, with int8 and int8-training modes for 8-bit inference and training, respectively, and NF4 or FP4 modes for 4-bit quantization. The QLoRA approach demonstrates that NF4 can support efficient fine-tuning of very large language models while retaining performance comparable to 16-bit fine-tuning, and its double-quantization option further reduces memory use by quantizing internal quantization constants through NF4-DQ or FP4-DQ modes.