Lightning-AI/pytorch-lightning

Pretrain, finetune ANY AI model of ANY size on 1 or 10,000+ GPUs with zero code changes.

What it solves

使用纯 PyTorch 训练深度学习模型通常需要编写重复且容易出错的工程代码，以处理基础设施层面的任务，例如反向传播、混合精度和跨多个 GPU 或节点的分布式训练。PyTorch Lightning 移除了这些样板代码，让研究人员能够专注于模型科学，而不是底层的基础设施工程。

How it works

Lightning 在 PyTorch 之上提供了两个层级的抽象化：

  1. PyTorch Lightning: 通过将模型逻辑（在 LightningModule 中定义）与训练循环（由 Trainer 处理）解耦，来组织 PyTorch 代码。这使得相同的代码可以在不改变核心模型逻辑的情况下，从 CPU 扩展到多节点 GPU 或 TPU。
  2. Lightning Fabric: 为复杂模型（如 LLMs 或基础模型）提供专家级控制。它提供了一种轻量级的方式来扩展 PyTorch 训练循环和策略（如 DDP, FSDP, 和 DeepSpeed），同时让用户保持对训练循环的控制权。

Who it’s for

想要在不管理分布式训练和硬件加速的复杂基础设施的情况下，预训练或微调模型（包括 LLMs, 扩散模型和图像分类器）的 AI 研究人员和开发人员。

Highlights

  • Hardware Agnostic: 通过简单的参数（flag）更改，即可在 CPU, GPU (CUDA/MPS), 和 TPU 之间切换，无需修改核心代码。
  • Massive Scalability: 支持在跨多个节点的数千个 GPU 上进行训练。
  • Built-in Optimizations: 原生支持 16-bit 混合精度和最先进的分布式策略，如 DeepSpeed 和 FSDP。
  • Extensive Integrations: 与 TensorBoard, Weights & Biases, Comet, 和 MLFlow 等热门的实验管理工具进行连接。
  • Production Ready: 包含用于将模型导出到 TorchScript (JIT) 和 ONNX 以进行部署的部署工具。